Inverse Burger PINN experiments (The automatic approach with PCGrad).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "be4c34c3-a418-4d45-8d18-9ea435524fb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from glob import glob as flist\n",
    "\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",
    "import pandas as pd\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 *\n",
    "\n",
    "# Multi-task learning loss, PCGrad\n",
    "import pcgrad\n",
    "\n",
    "# AdamGC (Gradient centrailization) optimizer\n",
    "# Please also try learning finder. (Doesn't have to be included in the paper)\n",
    "from optimizers import Lookahead, AdamGC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4d8aaa5f-7acc-4488-a7d7-9185217eb47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = \"./experimental_data/burgers_shock.mat\"\n",
    "\n",
    "data = io.loadmat(DATA_PATH)\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": "6b57ba04-7634-453e-aff7-3dd486298193",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.57870835*x0*x1 + 0.0021560264*x2\n"
     ]
    }
   ],
   "source": [
    "choice = 'manual'\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",
    "elif choice == 'manual':\n",
    "    # Grap the latestcsv file in the current directory\n",
    "    csv_file = sorted(flist('*csv'))[-1]\n",
    "    # Pick the equation that is the best (decided by pysr before) mannually\n",
    "    program = pd.read_csv(sorted(flist('*csv'))[-1], sep='|')['Equation'][3]\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)\n",
    "\n",
    "# Grap the latestcsv file in the current directory\n",
    "csv_file = sorted(flist('*csv'))[-1]\n",
    "# Pick the equation that is the best (decided by pysr before) mannually\n",
    "program = pd.read_csv(sorted(flist('*csv'))[-1], sep='|')['Equation'][3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0cfb62de-650d-42bb-80cf-06cd7ba315d9",
   "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": "30fed1f4-12ee-499c-a7eb-ead47d3e2b74",
   "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": "6ab19556-ebd4-4f05-9c1d-bfa0b0a8320f",
   "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.1395,  0.5787], 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": "2211934e-49f0-4698-93b8-45b8b9c0b216",
   "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 = []\n",
    "        uf = self.forward(x, t)\n",
    "        if update_network_params:\n",
    "            mse_loss = self.mean_squared(uf - y_input)\n",
    "            total_loss.append(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.append(l_eq)\n",
    "            \n",
    "        return total_loss\n",
    "    \n",
    "    def get_gradients_dict(self, x, t):\n",
    "        self.eval()\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": "9b14945c-430d-4f1b-b84c-39762f8abc3d",
   "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",
    "pretrained = False\n",
    "network = Network(model=model, pde_parameters=pde_params, loss_fn=fn, exp_indices=exp_indices)\n",
    "\n",
    "if pretrained:\n",
    "    print('Pretraining')\n",
    "    network.load_state_dict(torch.load(\"./saved_path_inverse_burger/nn_without_physical_reg.pth\"), strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1bc438a0-2711-40d9-82d9-e91adca3b528",
   "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": "5ecad091-ab46-4567-8681-aafcff729c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print('Mapping MSE loss')\n",
    "# ((network(X_star[:, 0:1], X_star[:, 1:2]).detach() - u_star)**2).mean().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "19cf7c9b-7f2e-408d-a5ee-dc3053e61b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 6e-3\n",
    "\n",
    "# optimizer1 = torch.optim.Adam(network.parameters(), lr=learning_rate)  # metaopt also has .parameters()\n",
    "optimizer1 = AdamGC(network.parameters(), lr=learning_rate, use_gc=True, gc_conv_only=False, gc_loc=False)\n",
    "epochs1 = 10000\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": 12,
   "id": "29cb8e52-fcbd-44d2-975c-5d3ba0214b00",
   "metadata": {},
   "outputs": [],
   "source": [
    "def closure():\n",
    "    if torch.is_grad_enabled():\n",
    "        optimizer2.zero_grad()\n",
    "    losses = network.loss(X_u_train[:, 0:1], X_u_train[:, 1:2], u_train, update_network_params=True, update_pde_params=True)\n",
    "    l = sum(losses)\n",
    "    if l.requires_grad:\n",
    "        l.backward()\n",
    "    return l\n",
    "\n",
    "def mtl_closure():\n",
    "    n_obj = 2 # There are two tasks\n",
    "    losses = network.loss(X_u_train[:, 0:1], X_u_train[:, 1:2], u_train, update_network_params=True, update_pde_params=True)\n",
    "    updated_grads = []\n",
    "    \n",
    "    for i in range(n_obj):\n",
    "        optimizer1.zero_grad()\n",
    "        losses[i].backward(retain_graph=True)\n",
    "\n",
    "        g_task = []\n",
    "        for param in network.parameters():\n",
    "            if param.grad is not None:\n",
    "                g_task.append(Variable(param.grad.clone(), requires_grad=False))\n",
    "            else:\n",
    "                g_task.append(Variable(torch.zeros(param.shape), requires_grad=False))\n",
    "        # appending the gradients from each task\n",
    "        updated_grads.append(g_task)\n",
    "\n",
    "    updated_grads = list(pcgrad.pc_grad_update(updated_grads))[0]\n",
    "    for idx, param in enumerate(network.parameters()): \n",
    "        param.grad = (updated_grads[0][idx]+updated_grads[1][idx])\n",
    "#         param.grad = (updated_grads[0][idx]+updated_grads[1][idx]).requires_grad_(True)\n",
    "        \n",
    "    return sum(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "60623efe-c5ff-4312-8abd-dc502e2f7068",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1st Phase optimization using Adam with PCGrad gradient modification\n",
      "Epoch 0:  0.6216228604316711\n",
      "Epoch 1000:  0.00715921213850379\n",
      "Epoch 2000:  0.0025596446357667446\n",
      "Epoch 3000:  0.0017578904516994953\n",
      "Epoch 4000:  0.0036378006916493177\n",
      "Epoch 5000:  0.0010350572410970926\n",
      "Epoch 6000:  0.0022573627065867186\n",
      "Epoch 7000:  0.036571841686964035\n",
      "Epoch 8000:  0.001264069927856326\n",
      "Epoch 9000:  0.00039816569187678397\n",
      "Epoch 9999:  0.009358081966638565\n",
      "2nd Phase optimization using LBFGS\n",
      "Epoch 0:  0.00710576307028532\n",
      "Epoch 100:  1.1807876944658346e-05\n",
      "Epoch 200:  7.445734809152782e-06\n",
      "Epoch 300:  7.445734809152782e-06\n",
      "Epoch 400:  7.445734809152782e-06\n",
      "Epoch 500:  7.445734809152782e-06\n",
      "Epoch 600:  7.445734809152782e-06\n",
      "Epoch 700:  7.445734809152782e-06\n",
      "Epoch 800:  7.445734809152782e-06\n",
      "Epoch 900:  7.445734809152782e-06\n",
      "Epoch 999:  7.445734809152782e-06\n"
     ]
    }
   ],
   "source": [
    "network.train(); best_train_loss = 1e6\n",
    "\n",
    "print('1st Phase optimization using Adam with PCGrad gradient modification')\n",
    "for i in range(epochs1):\n",
    "    optimizer1.step(mtl_closure)\n",
    "    l = mtl_closure()\n",
    "    \n",
    "    if (i % 1000) == 0 or i == epochs1-1:\n",
    "        print(\"Epoch {}: \".format(i), l.item())\n",
    "\n",
    "optimizer2 = torch.optim.LBFGS(network.parameters(), lr=5e-2, max_iter=50, max_eval=50, history_size=100, line_search_fn='strong_wolfe')\n",
    "epochs2 = 1000\n",
    "\n",
    "print('2nd Phase optimization using LBFGS')\n",
    "for i in range(epochs2):\n",
    "    optimizer2.step(closure)\n",
    "    l = closure()\n",
    "\n",
    "    if (i % 100) == 0 or i == epochs2-1:\n",
    "        print(\"Epoch {}: \".format(i), l.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b3fabdae-9750-4e47-8118-f9a59875b347",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0031862353649783444, 0.9994307160377502]"
      ]
     },
     "execution_count": 14,
     "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": 15,
   "id": "e2ed4f54-ac13-49a7-bf77-c03fd3768b71",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Estimating: [0.0031862353649783444, 0.9994307160377502]\n",
      "Real values: [0.003183098861837907, 1]\n",
      "Percentage error: 0.09853615223958938 0.056928396224975586\n",
      "Average percentage error 0.07773227423228249\n"
     ]
    }
   ],
   "source": [
    "### Loading the best weights ###\n",
    "# network.load_state_dict(torch.load(weights_path))\n",
    "\n",
    "network.eval()\n",
    "\n",
    "nu = 0.01 / np.pi\n",
    "if exp_indices[0] == 0: grounds = [nu, 1]\n",
    "else: grounds = [1, nu]\n",
    "\n",
    "print('Estimating:', est_lambdas)\n",
    "print('Real values:', grounds)\n",
    "\n",
    "error_lambda_1 = np.abs(est_lambdas[0] - grounds[0]) / grounds[0] * 100\n",
    "error_lambda_2 = np.abs(est_lambdas[1] - grounds[1]) / grounds[1] * 100\n",
    "\n",
    "print('Percentage error:', error_lambda_1, error_lambda_2)\n",
    "print('Average percentage error', 0.5*(error_lambda_1+error_lambda_2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ac270995-0218-4cca-a821-0470990d5b23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (0.11968308646106521 error on nu, 0.04165768623352051 error on 1)\n",
    "\n",
    "# with sparsity\n",
    "# Percentage error: 0.07810233039797228 on nu, 0.16725659370422363 on 1\n",
    "# Average percentage error 0.12267946205109795\n",
    "\n",
    "# without sparsity\n",
    "# Percentage error: 0.486716996366436 on nu, 0.025093555450439453 on 1\n",
    "# Average percentage error 0.2559052759084377"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}