Examples/Burgers-PMS-PINN-indecisive-ACS.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "53e1f133-f1d3-42c6-9c77-3615cb10d748",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sklearn's version: 1.2.1\n",
      "mrmr is not installed in the env you are using. This may cause an error in future if you try to use the (missing) lib.\n"
     ]
    }
   ],
   "source": [
    "# %load_ext autoreload\n",
    "# %autoreload 2\n",
    "# %matplotlib inline\n",
    "# import matplotlib.pyplot as plt\n",
    "import sys; sys.path.append('../')\n",
    "from misc import h5file\n",
    "\n",
    "import numpy as np\n",
    "from numpy.random import default_rng\n",
    "import scipy.io as sio\n",
    "\n",
    "import torch, sympytorch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import grad\n",
    "from siren_pytorch import SirenNet\n",
    "\n",
    "import pysindy as ps\n",
    "\n",
    "from sympy import symbols, simplify, lambdify\n",
    "from mathparser import math_eval\n",
    "from varname import nameof\n",
    "\n",
    "import sys; sys.path.append('../optimizers/')\n",
    "from Adan import Adan\n",
    "from FNO_Adam import Adam\n",
    "\n",
    "import sys; sys.path.append('../../parametric-discovery/')\n",
    "from tvregdiff import TVRegDiff, tvregdiff, numdiff, pysindydiff, savgol_denoise\n",
    "from functools import partial\n",
    "from best_subset import composite_function, ps_features\n",
    "import derivative\n",
    "\n",
    "class KalmanDiff(ps.BaseDifferentiation):\n",
    "    def __init__(self, alpha=1e-3, poly_deg=None, rpca_lambda=None, d=1, axis=1, is_uniform=True, periodic=False):\n",
    "        super(KalmanDiff, self).__init__()\n",
    "        # Kalman diff\n",
    "        self.alpha = alpha\n",
    "        self.diff_func = derivative.Kalman(alpha=self.alpha)\n",
    "        self.d = d\n",
    "        self.diff = partial(pysindydiff, **{\"diff_method\":self.diff_func, \"order\":self.d})\n",
    "        # Savgol denoising\n",
    "        self.poly_deg = poly_deg\n",
    "        if poly_deg is not None:\n",
    "            if poly_deg%2 == 0: window_length = self.poly_deg + 1\n",
    "            else: window_length = self.poly_deg + 2\n",
    "            self.denoise = partial(savgol_denoise, **{\"window_length\":window_length, \"poly_deg\":self.poly_deg})\n",
    "        else:\n",
    "            self.denoise = lambda _: _\n",
    "        # Robust PCA\n",
    "        self.rpca_lambda = rpca_lambda\n",
    "        # Other info...\n",
    "        self.axis = axis\n",
    "        self.is_uniform = is_uniform\n",
    "        self.periodic = periodic\n",
    "        # data transformation\n",
    "        # rs = np.ones(2).astype(np.int32); rs[self.axis] = -1; rs = tuple(rs)\n",
    "        self.transform = np.vectorize(composite_function(self.diff, self.denoise, left2right=True), signature=\"(m),(m)->(m)\")\n",
    "    # _differentiate\n",
    "    def _differentiate(self, x, t):\n",
    "        in_shape = x.shape\n",
    "        if len(in_shape) == 2: x = np.expand_dims(x, -1) # x should now be 3-dimensional\n",
    "        if isinstance(t, float) and self.is_uniform: \n",
    "            t = np.linspace(0, stop=t*(x.shape[self.axis]-1), num=x.shape[self.axis])\n",
    "        out = []\n",
    "        # wrt to x var\n",
    "        if self.axis == 0:\n",
    "            for i in range(x.shape[-1]):\n",
    "                # use lambda and partial from functools to help shorten the code\n",
    "                # diff = np.hstack([self.denoise(self.diff(x[:, j:j+1, i], t)).reshape(-1, 1) \n",
    "                #                   for j in range(x.shape[1])])\n",
    "                # diff = np.hstack([self.transform(x[:, j:j+1, i], t) for j in range(x.shape[1])])\n",
    "                # diff = np.vectorize(self.transform, signature=\"(m),(m)->(m)\")(x[:,:,i].T, t).T\n",
    "                diff = self.transform(x[:,:,i].T, t).T\n",
    "                if self.rpca_lambda is not None:\n",
    "                    diff = self._get_low_rank(diff)\n",
    "                out.append(np.expand_dims(diff, axis=-1))\n",
    "        # wrt to time var\n",
    "        elif self.axis == 1:\n",
    "            for i in range(x.shape[-1]):\n",
    "                # use lambda and partial from functools to help shorten the code\n",
    "                # diff = np.vstack([self.denoise(self.diff(x[j:j+1, :, i], t)).reshape(1, -1) \n",
    "                #                   for j in range(x.shape[0])])\n",
    "                # diff = np.vstack([self.transform(x[j:j+1, :, i], t) for j in range(x.shape[0])])\n",
    "                # diff = np.vectorize(self.transform, signature=\"(m),(m)->(m)\")(x[:,:,i], t)\n",
    "                diff = self.transform(x[:,:,i], t)\n",
    "                if self.rpca_lambda is not None:\n",
    "                    diff = self._get_low_rank(diff)\n",
    "                out.append(np.expand_dims(diff, axis=-1))\n",
    "        return np.concatenate(out, axis=-1).reshape(in_shape)\n",
    "    # _get_low_rank\n",
    "    def _get_low_rank(self, x):\n",
    "        rpca = RobustPCA(lamb=self.rpca_lambda, tol=10, use_fbpca=True, max_iter=int(1e6))\n",
    "        rpca.fit(x)\n",
    "        return rpca.get_low_rank()\n",
    "\n",
    "MAIN_SEED = 1234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a977d096-e837-4b20-98a7-32232f0866b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['X_pre', 'best_subsets', 'un', 'y_pre']\n"
     ]
    }
   ],
   "source": [
    "X_pre, best_subsets, un, y_pre = h5file(file_path=\"../PMS_data/indecisive_ACS/burgers_pms_indecisive_ACS.h5\", \n",
    "                                        mode='r', return_dict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c8deb99a-b9b4-47a0-afa1-58d7618f8416",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['u*u_1',\n",
       " 'u_11+u*u_1',\n",
       " 'u_11+u*u_1+u*u_11',\n",
       " 'u_1+u_11+u*u_1+u*u_11',\n",
       " 'u+u_11+u*u_1+u*u_11+u*u*u_11',\n",
       " 'u+u_11+u*u_1+u*u*u_1+u*u_11+u*u*u_11',\n",
       " 'u+u_1+u_11+u*u_1+u*u*u_1+u*u_11+u*u*u_11',\n",
       " 'u+u*u+u_1+u_11+u*u_1+u*u*u_1+u*u_11+u*u*u_11']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import yaml\n",
    "with open(\"../PMS_data/indecisive_ACS/burgers_pms_feature_names_indecisive_ACS.yaml\", 'r') as f:\n",
    "    config = yaml.load(f, yaml.Loader)\n",
    "f.close()\n",
    "encoded_feature_names = config[\"encoded_feature_names\"]\n",
    "encoded_pde_names = config[\"encoded_pde_names\"]\n",
    "encoded_pde_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "88a1b190-3187-419d-9da1-ca27b42430f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "poly_complexities = [name.count('*')+name.count('+')+1 for name in encoded_pde_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bcefd6e7-079d-458e-b952-891a52373d59",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = sio.loadmat('../Datasets/burgers.mat')\n",
    "x = (data['x'][0]).real\n",
    "t = (data['t'][:,0]).real\n",
    "dt = t[1]-t[0]; dx = x[2]-x[1]\n",
    "X, T = np.meshgrid(x, t)\n",
    "XT = np.asarray([X, T]).T\n",
    "del data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e290a14e-bff1-4d98-8824-7394dc44cf5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_like_value(prediction, ground):                                                                                                               \n",
    "    nobs = float(ground.shape[0])\n",
    "    nobs2 = nobs / 2.0\n",
    "    ssr = np.sum(np.abs(ground - prediction)**2)\n",
    "    llf = -nobs2 * np.log(2 * np.pi) - nobs2 * np.log(ssr / nobs) - nobs2\n",
    "    return llf\n",
    "\n",
    "def BIC_AIC(prediction, ground, nparams, reg_func = lambda x: x):\n",
    "    nparams = reg_func(nparams)\n",
    "    llf = log_like_value(prediction, ground)\n",
    "    return -2*llf + np.log(ground.shape[0])*nparams, -2*llf + 2*nparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "26cde66f-bb4c-4e52-ad4d-24e791b15669",
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_parameters(torch_model, onlyif_requires_grad=True):\n",
    "    if onlyif_requires_grad:\n",
    "        return sum(p.numel() for p in torch_model.parameters() if p.requires_grad)\n",
    "    return sum(p.numel() for p in torch_model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3193f8ba-74df-4b6a-ac18-bcdbb3a0935f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PhysicalConstraintCalculator(nn.Module):\n",
    "    def __init__(self, symbolic_module, basic_vars, init_coefficients=None, learnable_coefficients=False):\n",
    "        super(PhysicalConstraintCalculator, self).__init__()\n",
    "        self.symbolic_module = symbolic_module\n",
    "        self.basic_vars = basic_vars\n",
    "        \n",
    "        self.coefficients = init_coefficients\n",
    "        self.learnable_coefficients = learnable_coefficients\n",
    "\n",
    "        if self.coefficients is None:\n",
    "            self.coefficients = torch.ones(len(symbolic_module.sympy())).float()\n",
    "        else:\n",
    "            self.coefficients = torch.tensor(data=self.coefficients).float()\n",
    "        self.coefficients = nn.Parameter(self.coefficients).requires_grad_(self.learnable_coefficients)\n",
    "        \n",
    "        # printing\n",
    "        if self.learnable_coefficients: print(\"Learnable coefficients:\", self.coefficients)\n",
    "        else: print(\"NOT learnable coefficients:\", self.coefficients)\n",
    "        print(symbolic_module.sympy())\n",
    "        print(\"Basic variables:\", self.basic_vars)\n",
    "\n",
    "    def set_learnable_coefficients(self, learn):\n",
    "        self.coefficients.requires_grad_(learn)\n",
    "    \n",
    "    def forward(self, input_dict):\n",
    "        return self.symbolic_module(**input_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "27cf3237-bfb5-4215-8275-a22845c6b304",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sine(nn.Module):\n",
    "    def __init__(self, ):\n",
    "        super(Sine, self).__init__()\n",
    "    def forward(self, x):\n",
    "        return torch.sin(x)\n",
    "\n",
    "class TorchMLP(nn.Module):\n",
    "    def __init__(self, dimensions, bias=True, activation_function=nn.Tanh(), bn=None, dropout=None):\n",
    "        super(TorchMLP, self).__init__()\n",
    "        # setup ModuleList\n",
    "        self.model  = nn.ModuleList()\n",
    "        for i in range(len(dimensions)-1):\n",
    "            self.model.append(nn.Linear(dimensions[i], dimensions[i+1], bias=bias))\n",
    "            if bn is not None and i!=len(dimensions)-2:\n",
    "                self.model.append(bn(dimensions[i+1]))\n",
    "                if dropout is not None:\n",
    "                    self.model.append(dropout)\n",
    "            if i==len(dimensions)-2: break\n",
    "            self.model.append(activation_function)\n",
    "        # weight init\n",
    "        self.model.apply(self.xavier_init)\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.0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        for i, l in enumerate(self.model): \n",
    "            x = l(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bee2c43d-82d1-42d3-8696-0ffd0d23a6ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implement DeepONet(nn.Module)\n",
    "\n",
    "class PINN(nn.Module):\n",
    "    def __init__(self, solver, physics_calculator, lb, ub, \n",
    "                 domain_dimension=None, weak_pde_lib=None, effective_indices=None):\n",
    "        super(PINN, self).__init__()\n",
    "        self.solver = solver\n",
    "        self.physics_calculator = physics_calculator\n",
    "        self.lb = lb\n",
    "        self.ub = ub\n",
    "        # Only to use weak_loss\n",
    "        # spatial x temporal\n",
    "        self.domain_dimension = domain_dimension\n",
    "        self.weak_pde_lib = weak_pde_lib\n",
    "        self.effective_indices = effective_indices\n",
    "        self.weak_coeff_buffer = None\n",
    "        \n",
    "    def forward(self, x, t):\n",
    "        return self.solver(self.input_normalize(torch.cat([x, t],  dim=-1)))\n",
    "\n",
    "    def calculate_physics(self, x, t):\n",
    "        u = self.forward(x, t)\n",
    "        u_t = self.gradients(u, t)[0]\n",
    "        u_1 = self.gradients(u, x)[0]\n",
    "        u_11 = self.gradients(u_1, x)[0]\n",
    "        physics = self.physics_calculator({nameof(u):u, \n",
    "                                           nameof(u_1):u_1, \n",
    "                                           nameof(u_11):u_11})\n",
    "        \n",
    "        return u, u_t, physics\n",
    "    \n",
    "    def loss(self, x, t, y_input):\n",
    "        u, u_t, physics = self.calculate_physics(x, t)\n",
    "        coeff = self.physics_calculator.coefficients\n",
    "        physics = (physics*coeff).sum(axis=-1)\n",
    "        mse = F.mse_loss(u, y_input, reduction='mean')\n",
    "        l_eq = F.mse_loss(u_t, physics, reduction='mean')\n",
    "        return torch.add(mse, l_eq)\n",
    "    \n",
    "    def weak_loss(self, x, t, y_input):\n",
    "        u, u_t, physics = self.calculate_physics(x, t)\n",
    "        coeff = torch.tensor(self.weak_coefficients(u)).float()\n",
    "        physics = (physics*coeff).sum(axis=-1)\n",
    "        mse = F.mse_loss(u, y_input, reduction='mean')\n",
    "        l_eq = F.mse_loss(u_t, physics, reduction='mean')\n",
    "        return torch.add(mse, l_eq)\n",
    "    \n",
    "    def weak_form(self, u):\n",
    "        pred = u.reshape(self.domain_dimension[1], \n",
    "                         self.domain_dimension[0]).T.detach().numpy()\n",
    "        pred = np.expand_dims(pred,-1)\n",
    "        X_weak = self.weak_pde_lib.fit_transform(pred)\n",
    "        y_weak = self.weak_pde_lib.convert_u_dot_integral(pred)\n",
    "        return X_weak, y_weak\n",
    "    \n",
    "    def weak_coefficients(self, u):\n",
    "        np.random.seed(0)\n",
    "        X_weak, y_weak = self.weak_form(u)\n",
    "        X_weak = X_weak[:, self.effective_indices]\n",
    "        self.weak_coeff_buffer = np.linalg.lstsq(X_weak, y_weak, rcond=None)[0].flatten()\n",
    "        return self.weak_coeff_buffer\n",
    "    \n",
    "    def gradients(self, func, x):\n",
    "        return grad(func, x, create_graph=True, retain_graph=True, \n",
    "                    grad_outputs=torch.ones(func.shape))\n",
    "\n",
    "    def input_normalize(self, inp):\n",
    "        return -1.0+2.0*(inp-self.lb)/(self.ub-self.lb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d6cfb5d3-a870-4449-bf8c-b760a1d20373",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = default_rng(seed=0)\n",
    "# sampled_indices_x = rng.choice(len(x), size=int(len(x)//(2)), replace=False)\n",
    "# sampled_indices_t = rng.choice(len(t), size=int(len(t)//(2)), replace=False)\n",
    "sampled_indices_x = np.array([i for i in range(len(x)) if i%2==0])\n",
    "sampled_indices_t = np.array([i for i in range(len(t)) if i%2==0]) \n",
    "domain_dimension = len(sampled_indices_x), len(sampled_indices_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1be26291-1b4d-4cf9-9d41-dfce3c1913a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(MAIN_SEED);\n",
    "torch.manual_seed(MAIN_SEED);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c632efe7-a59b-4ef1-9939-fb1037932e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "XX = X[sampled_indices_t, :][:, sampled_indices_x]\n",
    "TT = T[sampled_indices_t, :][:, sampled_indices_x]\n",
    "XXTT = XT[sampled_indices_x, :, :][:, sampled_indices_t, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8173a691-9bfb-406b-ab0a-4a3ab27852c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 3000; diff_order = 2\n",
    "weak_pde_lib = ps.WeakPDELibrary(library_functions=[lambda x:x, lambda x: x*x], \n",
    "                                 function_names=[lambda x:x, lambda x: x+x], \n",
    "                                 derivative_order=diff_order, p=diff_order, \n",
    "                                 spatiotemporal_grid=XXTT, \n",
    "                                 include_bias=False, is_uniform=True, K=K # new random K points in every calls to the ps.WeakPDELibrary\n",
    "                                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "99bd1108-8cab-4908-9f84-23226465c355",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = np.hstack((XX.flatten()[:,None], TT.flatten()[:,None]))\n",
    "y_train = un.T[sampled_indices_t, :][:, sampled_indices_x].flatten()[:,None]\n",
    "lb = torch.tensor(X_train.min(axis=0)).float().requires_grad_(False)\n",
    "ub = torch.tensor(X_train.max(axis=0)).float().requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5ac6a7c0-6773-4b9a-b6cd-382d767f4385",
   "metadata": {},
   "outputs": [],
   "source": [
    "del XX, TT, XXTT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cbff35fc-258f-4bb9-b71c-543237cbb3d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([6528, 2]), torch.Size([6528, 1]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Converting to tensors\n",
    "X_train = torch.tensor(X_train).float().requires_grad_(True)\n",
    "y_train = torch.tensor(y_train).float().requires_grad_(False)\n",
    "X_train.shape, y_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f745a25b-5bca-45fd-ac3a-7ee9010277fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 0.09748721, -1.0107358 ], dtype=float32),\n",
       " SymPyModule(expressions=(u_11, u*u_1)))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "com = 2; com = max(com, 1)\n",
    "effective_indices = np.where(best_subsets[com-1]>0)[0].tolist()\n",
    "init_coefficients = np.linalg.lstsq(X_pre[:, effective_indices], \n",
    "                                    y_pre, rcond=None)[0].flatten()\n",
    "mod, basic_vars = math_eval(encoded_pde_names[com-1], \n",
    "                            return_torch=True, split_by_addition=True)\n",
    "init_coefficients, mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "70d17a5a-5c79-4012-aa86-4a44dee55eba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learnable coefficients: Parameter containing:\n",
      "tensor([ 0.0975, -1.0107], requires_grad=True)\n",
      "[u_11, u*u_1]\n",
      "Basic variables: ['u', 'u_1', 'u_11']\n"
     ]
    }
   ],
   "source": [
    "# bias init at 0.01 | SIREN\n",
    "# activation_function = nn.Tanh()\n",
    "activation_function = Sine()\n",
    "n_nodes = 10 # tuned by min BIC on train\n",
    "solver = TorchMLP([2,n_nodes,n_nodes,n_nodes,n_nodes,1], \n",
    "                  bn=None, activation_function=activation_function)\n",
    "# solver = SirenNet(dim_in=2, dim_hidden=50, dim_out=1, num_layers = 4, \n",
    "#                   w0_initial = 30.)\n",
    "physics_calculator = PhysicalConstraintCalculator(symbolic_module=mod, \n",
    "                                                  basic_vars=basic_vars, \n",
    "                                                  init_coefficients=init_coefficients, \n",
    "                                                  learnable_coefficients=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4f2530f7-851b-4c3d-b42b-6b4f61ef432c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pinn = PINN(solver, physics_calculator, \n",
    "            lb, ub, domain_dimension, \n",
    "            weak_pde_lib, effective_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a354caae-01c5-40fb-a672-2d861cb62cf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "# pinn.weak_loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "# count_parameters(pinn, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "239324b8-04d0-4c9e-af66-5c77c05876a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.005399589892476797\n",
      "Epoch 50:  0.005229737143963575\n",
      "Epoch 99:  0.005229737143963575\n"
     ]
    }
   ],
   "source": [
    "pinn.physics_calculator.set_learnable_coefficients(False)\n",
    "lbfgs = torch.optim.LBFGS(pinn.parameters(), lr=0.1, \n",
    "                          max_iter=500, max_eval=500, history_size=300, \n",
    "                          line_search_fn='strong_wolfe')\n",
    "epochs = 100\n",
    "pinn.train()\n",
    "\n",
    "for i in range(epochs):\n",
    "    def closure():\n",
    "        if torch.is_grad_enabled(): \n",
    "            lbfgs.zero_grad()\n",
    "        l = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "        if l.requires_grad: \n",
    "            l.backward()\n",
    "        return l\n",
    "\n",
    "    lbfgs.step(closure)\n",
    "\n",
    "    # calculate the loss again for monitoring\n",
    "    if (i%50)==0 or i==epochs-1:\n",
    "        l = closure()\n",
    "        print(\"Epoch {}: \".format(i), l.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b2883872-bd37-4aaf-aed2-920c43b59c6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_base: 371\n",
      "(-15786.583581447941, -15800.15129324123)\n",
      "(-12527.77304379304, -15058.15129324123)\n",
      "(-12518.989187896395, -15056.15129324123)\n"
     ]
    }
   ],
   "source": [
    "### Indecisive ACS ### -> BIC is better!!! (more regularization)\n",
    "### params -> base = 7851\n",
    "### com=1 ###\n",
    "# train\n",
    "# (-15654.985790517178, -15661.76964641382)\n",
    "# (53307.06685403441, 40.23035358617926)\n",
    "# (53315.85070993106, 42.23035358617926)\n",
    "# val\n",
    "# (-15383.198594154986, -15389.962647424334)\n",
    "# (53423.383623494294, 312.03735257566586)\n",
    "# (53432.14767676364, 314.03735257566586)\n",
    "# ----------\n",
    "### com=2 ###\n",
    "# train\n",
    "# (-15767.549792022857, -15781.117503816145)\n",
    "# (53194.502852528734, -79.11750381614547)\n",
    "# (53203.28670842538, -77.11750381614547)\n",
    "# val\n",
    "# (-15518.561023892013, -15532.089130430708)\n",
    "# (53288.02119375727, 169.91086956929212) \n",
    "# (53296.785247026615, 171.91086956929212) ***choose***\n",
    "# ----------\n",
    "### com=3 ###\n",
    "# train\n",
    "# (-15763.747879830795, -15784.099447520726)\n",
    "# (53198.3047647208, -82.09944752072624)\n",
    "# (53215.87247651408, -78.09944752072624)\n",
    "# val\n",
    "# (-15512.137580354043, -15532.429740162086)\n",
    "# (53294.44463729524, 169.5702598379139)\n",
    "# (53311.97274383393, 173.5702598379139)\n",
    "\n",
    "### n_base: 1341\n",
    "### com = 2 ###\n",
    "# train\n",
    "# (-15774.099529915524, -15787.667241708812)\n",
    "# (-3994.948772515998, -13105.667241708812)\n",
    "# (-3986.164916619353, -13103.667241708812)\n",
    "# val\n",
    "# (-15517.064811271412, -15530.592917810107)\n",
    "# (-3764.4693770760623, -12848.592917810107)\n",
    "# (-3755.7053238067147, -12846.592917810107)\n",
    "# ----------\n",
    "### com = 3 ###\n",
    "# train\n",
    "# (-15776.894748991685, -15797.246316681616)\n",
    "# (-3997.743991592157, -13115.246316681616)\n",
    "# (-3980.1762797988704, -13111.246316681616)\n",
    "# val\n",
    "# (-15503.868795770904, -15524.160955578947)\n",
    "# (-3751.273361575555, -12842.160955578947)\n",
    "# (-3733.74525503686, -12838.160955578947)\n",
    "\n",
    "### n_base: 371 ###\n",
    "### com=2 ###\n",
    "# train\n",
    "# (-15786.583581447941, -15800.15129324123)\n",
    "# (-12527.77304379304, -15058.15129324123)\n",
    "# (-12518.989187896395, -15056.15129324123)\n",
    "# val\n",
    "# (-15506.363844480125, -15519.89195101882)\n",
    "# (-12254.900081552105, -14777.89195101882)\n",
    "# (-12246.136028282757, -14775.89195101882)\n",
    "# ----------\n",
    "### com=3 ###\n",
    "# train\n",
    "# (-15759.435826522693, -15779.787394212624)\n",
    "# (-12500.62528886779, -15037.787394212624)\n",
    "# (-12483.057577074502, -15033.787394212624)\n",
    "# val\n",
    "# (-15501.242202374975, -15521.534362183018)\n",
    "# (-12249.778439446955, -14779.534362183018)\n",
    "# (-12232.25033290826, -14775.534362183018)\n",
    "\n",
    "pinn.eval()\n",
    "pred = pinn(X_train[:, 0:1], X_train[:, 1:2]).detach().numpy()\n",
    "base = count_parameters(pinn, False)-com; print(\"n_base:\", base)\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), com))\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), count_parameters(pinn, False)))\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), base+poly_complexities[com-1])) # a good choice with AIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ae4d8620-a917-4f3c-b236-1a014e785015",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(pinn.state_dict(), \"./tmp_files/tmp1.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "517c22b5-4b81-416b-bdb2-4ed21baa47e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_indices_x = np.array([i for i in range(len(x)) if i%2==1])\n",
    "validation_indices_t = np.array([i for i in range(len(t)) if i%2==1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "12a8bcd2-c8f1-4cc5-a534-5bfbd594d199",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(-15506.363844480125, -15519.89195101882)\n",
      "(-12254.900081552105, -14777.89195101882)\n",
      "(-12246.136028282757, -14775.89195101882)\n"
     ]
    }
   ],
   "source": [
    "val_pred = pinn(torch.tensor(X[validation_indices_t, :][:, validation_indices_x].flatten()[:,None]).float(), \n",
    "                torch.tensor(T[validation_indices_t, :][:, validation_indices_x].flatten()[:,None]).float()).detach().numpy()\n",
    "y_val = un.T[validation_indices_t, :][:, validation_indices_x].flatten()[:,None]\n",
    "print(BIC_AIC(val_pred, y_val, com))\n",
    "print(BIC_AIC(val_pred, y_val, count_parameters(pinn, False)))\n",
    "print(BIC_AIC(val_pred, y_val, base+poly_complexities[com-1])) # a good choice with AIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6764eb48-0d8a-4786-9c2c-fc6fbb72397c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "512810ac-c53a-48db-883b-1f7a49097f1e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "22a3e27b-7dff-4f8c-aff7-da75eac15782",
   "metadata": {},
   "source": [
    "#### weak"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "40590930-a24c-407a-a0b4-81a00b7fe862",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### using lbfgs ###\n",
    "# print(\"using lbfgs...\")\n",
    "# epochs = 100\n",
    "# pinn.train()\n",
    "# pinn.physics_calculator.set_learnable_coefficients(False)\n",
    "# lbfgs2 = torch.optim.LBFGS(pinn.parameters(), lr=0.1, \n",
    "#                           max_iter=500, max_eval=500, history_size=300, \n",
    "#                           line_search_fn='strong_wolfe')\n",
    "\n",
    "# for i in range(epochs):\n",
    "#     def closure2():\n",
    "#         if torch.is_grad_enabled(): \n",
    "#             lbfgs2.zero_grad()\n",
    "#         l = pinn.weak_loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "#         if l.requires_grad: \n",
    "#             l.backward()\n",
    "#         return l\n",
    "\n",
    "#     lbfgs2.step(closure2)\n",
    "\n",
    "#     # calculate the loss again for monitoring\n",
    "#     if (i%50)==0 or i==epochs-1:\n",
    "#         l = closure2()\n",
    "#         print(\"Epoch {}: \".format(i), l.item())\n",
    "\n",
    "# pinn.eval()\n",
    "# pred = pinn(X_train[:, 0:1], X_train[:, 1:2]).detach().numpy()\n",
    "# print(BIC_AIC(pred, y_train.detach().numpy(), com))\n",
    "# print(pinn.physics_calculator.coefficients.detach().numpy())\n",
    "# print(pinn.weak_coeff_buffer)\n",
    "\n",
    "# ### using non-lbfgs ###\n",
    "# print(\"using non-lbfgs...\")\n",
    "# epochs = 1000\n",
    "# pinn.train()\n",
    "# pinn.physics_calculator.set_learnable_coefficients(False)\n",
    "# # optimizer = Adam(pinn.parameters(), lr=1e-5, weight_decay=1e-4)\n",
    "# optimizer = Adan(pinn.parameters(), lr=1e-5, weight_decay=1e-2)\n",
    "\n",
    "# for i in range(epochs):\n",
    "#     optimizer.zero_grad()\n",
    "#     l = pinn.weak_loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "#     l.backward()\n",
    "#     optimizer.step()\n",
    "#     if (i%50)==0 or i==epochs-1:\n",
    "#         l = pinn.weak_loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "#         print(\"Epoch {}: \".format(i), l.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2a2b635-4091-4721-8b23-eaddaab8d080",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab22b1c-5010-49e8-89ff-d58c205806cd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fed5db4c-235e-42ad-9e94-7f8ee1ad0a2c",
   "metadata": {},
   "source": [
    "#### learn (pinn.physics_calculator.set_learnable_coefficients(True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d3415537-5acf-4f6b-9e73-8166a1c73be3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### using lbfgs ###\n",
    "# epochs = 100\n",
    "# pinn.train()\n",
    "# pinn.physics_calculator.set_learnable_coefficients(True)\n",
    "# lbfgs2 = torch.optim.LBFGS(pinn.parameters(), lr=0.1, \n",
    "#                           max_iter=500, max_eval=500, history_size=300, \n",
    "#                           line_search_fn='strong_wolfe')\n",
    "\n",
    "# for i in range(epochs):\n",
    "#     def closure2():\n",
    "#         if torch.is_grad_enabled(): \n",
    "#             lbfgs2.zero_grad()\n",
    "#         l = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "#         if l.requires_grad: \n",
    "#             l.backward()\n",
    "#         return l\n",
    "\n",
    "#     lbfgs2.step(closure2)\n",
    "\n",
    "#     # calculate the loss again for monitoring\n",
    "#     if (i%50)==0 or i==epochs-1:\n",
    "#         l = closure2()\n",
    "#         print(\"Epoch {}: \".format(i), l.item())\n",
    "        \n",
    "# ### using non-lbfgs ###\n",
    "# epochs = 1000\n",
    "# pinn.train()\n",
    "# pinn.physics_calculator.set_learnable_coefficients(True)\n",
    "# # optimizer = Adam(pinn.parameters(), lr=1e-5, weight_decay=1e-4)\n",
    "# optimizer = Adan(pinn.parameters(), lr=1e-5, weight_decay=1e-2)\n",
    "\n",
    "# for i in range(epochs):\n",
    "#     optimizer.zero_grad()\n",
    "#     l = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "#     l.backward()\n",
    "#     optimizer.step()\n",
    "#     if (i%50)==0 or i==epochs-1:\n",
    "#         l = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "#         print(\"Epoch {}: \".format(i), l.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b761000-0dfa-472b-a07c-6d75d51afba7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a582100c-c27f-4cc3-bf31-48de02e21c67",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b5b40b39-5aa5-44d1-88b9-cc47fa93c2c1",
   "metadata": {},
   "source": [
    "#### eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "fdd38d55-1674-4c54-89b5-6e54361dc6c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(-15786.583581447941, -15800.15129324123)\n",
      "[ 0.09748721 -1.0107358 ]\n"
     ]
    }
   ],
   "source": [
    "pinn.eval()\n",
    "pred = pinn(X_train[:, 0:1], X_train[:, 1:2]).detach().numpy()\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), com))\n",
    "print(pinn.physics_calculator.coefficients.detach().numpy())\n",
    "# print(pinn.weak_coeff_buffer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8df2aaca-0473-42a0-9cc9-6a5b809d935a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1.7931849999999985, 0.7196050000000023)\n",
      "(1.7931818962097195, 0.7196068763732937)\n"
     ]
    }
   ],
   "source": [
    "def percent_coeff_error(pred):\n",
    "    ground = np.array([0.1, -1])\n",
    "    errs = 100*np.abs(np.array(pred)-ground)/np.abs(ground)\n",
    "    return errs.mean(), errs.std()\n",
    "print(percent_coeff_error([ 0.09748721, -1.0107358 ]))\n",
    "print(percent_coeff_error(pinn.physics_calculator.coefficients.detach().numpy().tolist()))\n",
    "# print(percent_coeff_error(pinn.weak_coeff_buffer.tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "79a5a3d9-7cfc-47c6-b8e8-883816d68aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_domain_pred = pinn(torch.tensor(X.flatten()[:,None]).float(), \n",
    "                        torch.tensor(T.flatten()[:,None]).float()).detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "6189fb63-c004-4c39-ad96-c50671c3a4e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# differentiation_method = ps.FiniteDifference\n",
    "# differentiation_kwargs = {}\n",
    "kalpha = 1e-3; poly_deg = None\n",
    "differentiation_method = KalmanDiff\n",
    "differentiation_kwargs = {\"alpha\":kalpha, \"poly_deg\":poly_deg, \"rpca_lambda\":None}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "404d4520-8dc0-46b1-bd68-9103c4eb6956",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_mean = np.zeros(X_pre.shape)\n",
    "y_mean = np.zeros(y_pre.shape)\n",
    "n_times = 10\n",
    "np.random.seed(0)\n",
    "for _ in range(n_times):\n",
    "    weak_kalman_pde_lib = ps.WeakPDELibrary(library_functions=[lambda x:x, lambda x: x*x], \n",
    "                                            function_names=[lambda x:x, lambda x: x+x], \n",
    "                                            derivative_order=2, p=2, \n",
    "                                            spatiotemporal_grid=XT, \n",
    "                                            include_bias=False, is_uniform=True, K=X_mean.shape[0], \n",
    "                                            differentiation_method=differentiation_method, \n",
    "                                            differentiation_kwargs=differentiation_kwargs, \n",
    "                                            cache=False\n",
    "                                           )\n",
    "    kwargs = {'fit_intercept':False, 'copy_X':True, 'normalize_columns':False}\n",
    "    X_mean_sub, y_mean_sub, _ = ps_features(full_domain_pred.reshape(len(t), len(x)).T, \n",
    "                                            t, weak_kalman_pde_lib, kwargs)\n",
    "    X_mean = X_mean + X_mean_sub\n",
    "    y_mean = y_mean + y_mean_sub\n",
    "    \n",
    "X_mean = X_mean/n_times\n",
    "y_mean = y_mean/n_times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "15b5105c-b12b-489c-b54f-216903b018ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.10073452 -0.99778277]\n",
      "(0.4781207374699059, 0.2563977082039426)\n"
     ]
    }
   ],
   "source": [
    "final_coeff = np.linalg.lstsq(X_mean[:, [3,4]], y_mean, rcond=None)[0].flatten()\n",
    "print(final_coeff)\n",
    "print(percent_coeff_error(final_coeff))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "3d089b58-af3d-4a19-8d28-520b62904115",
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.api import OLS as SMOLS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "4bca5614-e624-4066-9826-b1e562183f49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.10073452 -0.99778277]\n",
      "(0.4781207374699087, 0.25639770820391206)\n",
      "-33886.2806497275\n"
     ]
    }
   ],
   "source": [
    "mr = SMOLS(y_mean, X_mean[:, np.where(best_subsets[0]>0)[0].tolist()]).fit()\n",
    "mb = SMOLS(y_mean, X_mean[:, np.where(best_subsets[1]>0)[0].tolist()]).fit()\n",
    "met = (mb.bic-mr.bic)/(len(mb.params)-len(mr.params))\n",
    "print(mb.params)\n",
    "print(percent_coeff_error(mb.params))\n",
    "print(met)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "d6a8a30f-65fa-4ca0-bbb4-401469b7ab9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ps.FiniteDiff\n",
    "# [ 0.10123641 -1.00077165]\n",
    "# (0.6567893638010344, 0.5796248482694429)\n",
    "\n",
    "# 1e-1 | -33848.877106793036\n",
    "# [ 0.10555368 -1.00447551]\n",
    "# (3.0006146867824404, 2.5530640898926507)\n",
    "# 1e-3 | -33886.2806497275\n",
    "# (0.4781207374699087, 0.25639770820391206)\n",
    "# 1e-6 | -33886.050614698644\n",
    "# [ 0.10073346 -0.99778198]\n",
    "# (0.47762903455395994, 0.2558274380621048)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e92c22a8-6da9-4470-b9b7-f1580ea9ec59",
   "metadata": {},
   "source": [
    "#### Notes\n",
    "    - 3 main files are required."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pysr]",
   "language": "python",
   "name": "conda-env-pysr-py"
  },
  "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.10.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}