Examples/Burgers-PMS-PINN-noise30-IC-pysr.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.2\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",
    "from itertools import combinations\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "from numpy.random import default_rng\n",
    "import scipy.io as sio\n",
    "from scipy.optimize import curve_fit\n",
    "from jaxfit import CurveFit\n",
    "from levenberg_marquardt import lm as lm_curve_fit\n",
    "from statsmodels.api import OLS as SMOLS\n",
    "import sympy\n",
    "import pandas as pd\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, sympify, 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",
    "\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",
    "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",
    "\n",
    "MAIN_SEED = 1234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3f7949ae-1fe1-4c45-a36b-03c50947e4e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class KalmanDiff(ps.BaseDifferentiation):\n",
    "    def __init__(self, alpha, 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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f8e180f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('../PMS_data/ksvd/none/transform_n_nonzero_coefs_none/burgers_pms_noise30_dictlearn32.h5',\n",
       " '../PMS_data/ksvd/none/transform_n_nonzero_coefs_none/burgers_pms_feature_names_noise30_dictlearn32.yaml')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "noise_lv = 30\n",
    "denoising_mode = 'ksvd'\n",
    "smoother_name = 'none'\n",
    "n_components = 32\n",
    "transform_n_nonzero = '_none'\n",
    "undenoised = False\n",
    "\n",
    "fp1 = f\"../PMS_data/{denoising_mode}/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/burgers_pms_noise30_dictlearn{n_components}.h5\"\n",
    "fp2 = f\"../PMS_data/{denoising_mode}/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/burgers_pms_feature_names_noise30_dictlearn{n_components}.yaml\"\n",
    "\n",
    "fp1, fp2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a977d096-e837-4b20-98a7-32232f0866b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['X_pre', 'avg_weak_coeff', 'best_subsets', 'un', 'y_pre']\n"
     ]
    }
   ],
   "source": [
    "# RDAE, noRDAE\n",
    "X_pre, avg_weak_coeff, best_subsets, un, y_pre = \\\n",
    "h5file(file_path=fp1, mode='r', return_dict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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*u+u_11+u*u_1+u*u_11',\n",
       " 'u*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*u+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": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import yaml\n",
    "# RDAE, noRDAE\n",
    "with open(fp2, '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": 6,
   "id": "58ed3b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [np.linalg.lstsq(X_pre[:, np.where(best_subsets[i]>0)[0]], \n",
    "#                  y_pre, rcond=None)[0].flatten() for i in range(len(best_subsets))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 8,
   "id": "bcefd6e7-079d-458e-b952-891a52373d59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In the Denoised mode...\n"
     ]
    }
   ],
   "source": [
    "data = sio.loadmat('../Datasets/burgers.mat')\n",
    "\n",
    "u_clean = data['usol'].real\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",
    "\n",
    "if undenoised:\n",
    "    np.random.seed(0)\n",
    "    un = u_clean + 0.01*np.abs(noise_lv)*(u_clean.std())*np.random.randn(u_clean.shape[0], \n",
    "                                                                         u_clean.shape[1])\n",
    "    print(\"In the Undenoised mode...\")\n",
    "else:\n",
    "    print(\"In the Denoised mode...\")\n",
    "    \n",
    "# del data, u_clean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 10,
   "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": 11,
   "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": 12,
   "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": 13,
   "id": "e7558a8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "                 ic_module=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",
    "        # must not be None if X_train_initial is not None but y_train_initial is None\n",
    "        self.ic_module = ic_module\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, X_train_initial=None, y_train_initial=None):\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",
    "        \n",
    "        # initial condition (ic)\n",
    "        if X_train_initial is not None:\n",
    "            ic_u_pred = self.solver(self.input_normalize(X_train_initial))\n",
    "            if y_train_initial is None:\n",
    "                y_train_initial = self.ic_module(X_train_initial)\n",
    "            ic_loss = F.mse_loss(ic_u_pred, y_train_initial, reduction='mean')\n",
    "            mse = torch.add(mse, ic_loss)\n",
    "            \n",
    "        l_eq = F.mse_loss(u_t, physics, reduction='mean')\n",
    "        return 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",
    "        \n",
    "        # initial condition (ic)\n",
    "        if X_train_initial is not None:\n",
    "            ic_u_pred = self.solver(self.input_normalize(X_train_initial))\n",
    "            if y_train_initial is None:\n",
    "                y_train_initial = self.ic_module(X_train_initial)\n",
    "            ic_loss = F.mse_loss(ic_u_pred, y_train_initial, reduction='mean')\n",
    "            mse = torch.add(mse, ic_loss)\n",
    "            \n",
    "        l_eq = F.mse_loss(u_t, physics, reduction='mean')\n",
    "        return 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 set_learnable_ic(self, flag):\n",
    "        self.ic_module.requires_grad_(flag)\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": 14,
   "id": "d6cfb5d3-a870-4449-bf8c-b760a1d20373",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finetuning\n"
     ]
    }
   ],
   "source": [
    "rng = default_rng(seed=0)\n",
    "mode = 'finetuning' # 'selection', finetuning'\n",
    "if mode == 'finetuning':\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",
    "elif mode == 'selection':\n",
    "    sampled_indices_x = np.array([i for i in range(len(x)) if i<len(x)//2+1])\n",
    "#     sampled_indices_t = np.array([i for i in range(len(t)) if i<len(t)//2+1 and i!=0])\n",
    "    sampled_indices_t = np.array([i for i in range(len(t)) if i<len(t)//2+1])\n",
    "print(mode)\n",
    "domain_dimension = len(sampled_indices_x), len(sampled_indices_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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": 16,
   "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": 17,
   "id": "8173a691-9bfb-406b-ab0a-4a3ab27852c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# kalpha = 1e-1; poly_deg = None\n",
    "# differentiation_method = KalmanDiff\n",
    "# differentiation_kwargs = {\"alpha\":kalpha, \"poly_deg\":poly_deg, \"rpca_lambda\":None}\n",
    "\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",
    "#                                  differentiation_method=differentiation_method, \n",
    "#                                  differentiation_kwargs=differentiation_kwargs, \n",
    "#                                  cache=False\n",
    "#                                 )\n",
    "\n",
    "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": 18,
   "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.min(), t.min()]).float().requires_grad_(False)\n",
    "# ub = torch.tensor([x.max(), t.max()]).float().requires_grad_(False)\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": 19,
   "id": "d8ce57ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "del XX, TT, XXTT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6f259f8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exp(-1.0023574 * square(2.0152135 + x0))\n"
     ]
    }
   ],
   "source": [
    "# Actually this eq is valid for smoother_name = 'kalman' only but anyways...\n",
    "hof = pd.read_csv(\"./hof.csv\")\n",
    "if smoother_name != 'none': \n",
    "    hof = pd.read_csv(f\"./hof_{smoother_name}.csv\")\n",
    "equation = hof.iloc[np.argmax(hof[\"score\"])]\n",
    "# how to extract float numbers from a sympy object?\n",
    "print(equation.equation)\n",
    "func = lambdify(args=sympy.symbols('x0'), expr=equation.equation)\n",
    "pysr_params = np.array(sorted([float(atom) for atom in sympify(equation.equation).atoms() if atom.is_number]))\n",
    "initial_indices = np.where(X_train[:, 1:2]==0.0)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b9aa65db",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-17 22:10:09,756 [INFO] Remote TPU is not linked into jax; skipping remote TPU.\n",
      "2023-04-17 22:10:09,757 [INFO] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'\n",
      "2023-04-17 22:10:09,758 [INFO] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2023-04-17 22:10:09,759 [INFO] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2023-04-17 22:10:09,760 [INFO] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'\n",
      "2023-04-17 22:10:09,762 [INFO] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using uniform weights for error analysis\n",
      "**** Convergence in r.h.s. (\"JtWdy\")  ****\n",
      "\n",
      "LM fitting results:\n",
      "----------------------------- \n",
      "parameter      = p1\n",
      "fitted value   = -1.0027\n",
      "standard error = -2.38 %\n",
      "----------------------------- \n",
      "parameter      = p2\n",
      "fitted value   = 2.0155\n",
      "standard error = 1.19 %\n",
      "[-1.0023574  2.0152135]\n"
     ]
    }
   ],
   "source": [
    "def initial_function(x, a, b): return np.exp(a*np.square(x+b))\n",
    "def jax_initial_function(x, a, b): return jnp.exp(a*jnp.square(x+b))\n",
    "\n",
    "recovered_params1 = np.array(CurveFit().curve_fit(jax_initial_function, x.flatten(), un[:, 0], \n",
    "                                                  p0=np.round(pysr_params))[0]) # p0=np.round(pysr_params), p0=None\n",
    "\n",
    "recovered_params2 = np.array(curve_fit(initial_function, x.flatten(), un[:, 0], \n",
    "                                       p0=None, method='lm')[0])\n",
    "\n",
    "recovered_params3 = lm_curve_fit(np.round(pysr_params).reshape(-1, 1), \n",
    "                                 x.flatten(), un[:, 0], \n",
    "                                 lambda t,p: np.exp(p[0,0]*np.square(t+p[1,0])))[0].flatten()\n",
    "\n",
    "# pysr_params, recovered_params1, recovered_params2, recovered_params3 (recommended when finetuning)\n",
    "recovered_params = pysr_params\n",
    "# recovered_params = np.round(recovered_params, decimals=6)\n",
    "print(recovered_params)\n",
    "\n",
    "initial_function = partial(initial_function, a=recovered_params[0], b=recovered_params[1])\n",
    "# initial_function = partial(initial_function, a=-1.0, b=2.0) # GROUND\n",
    "# initial_function = func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0dc9d3be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'_ic'"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train_initial, y_train_initial = None, None\n",
    "add_initial_data = 2 # 0, 1, 2\n",
    "\n",
    "X0, T0 = np.meshgrid(x, np.array([0.0]))\n",
    "\n",
    "if add_initial_data == 1:\n",
    "    ### V1 of adding initial data ###\n",
    "    if len(initial_indices) > 0:\n",
    "        y_train[initial_indices] = np.vectorize(initial_function)(X_train[initial_indices][:, 0:1])\n",
    "elif add_initial_data == 2:\n",
    "    ### V2 of adding initial data ###\n",
    "    if add_initial_data:\n",
    "        X_train_initial = np.hstack((X0.flatten()[:,None], T0.flatten()[:,None]))\n",
    "        y_train_initial = initial_function(X_train_initial[:, 0:1])\n",
    "        X_train_initial = torch.tensor(X_train_initial).float().requires_grad_(False)\n",
    "        y_train_initial = torch.tensor(y_train_initial).float().requires_grad_(False)\n",
    "\n",
    "if add_initial_data>0:\n",
    "    if np.abs(recovered_params-pysr_params).sum() == 0.0: \n",
    "        with_initial_data = '_ic'\n",
    "    else: \n",
    "        with_initial_data = '_lm'\n",
    "else:\n",
    "     with_initial_data = ''\n",
    "        \n",
    "del X0, T0\n",
    "\n",
    "with_initial_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b6f123a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([6528, 2]), torch.Size([6528, 1]))"
      ]
     },
     "execution_count": 23,
     "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": 24,
   "id": "f745a25b-5bca-45fd-ac3a-7ee9010277fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 0.10379176, -1.00721217]), SymPyModule(expressions=(u_11, u*u_1)))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_com = 2\n",
    "com = 2; com = max(com, 1)\n",
    "\n",
    "# getting effective_indices\n",
    "# effective_indices = np.where(best_subsets[com-1]>0)[0].tolist()\n",
    "all_subsets = list(combinations(range(len(config[\"encoded_feature_names\"])), com))\n",
    "scores = []\n",
    "for s in all_subsets:\n",
    "    inp = X_pre[:, s]\n",
    "    w = np.linalg.lstsq(inp, y_pre, rcond=None)[0]\n",
    "    scores.append(((y_pre-inp@w)**2).mean())\n",
    "effective_indices = all_subsets[np.argmin(scores)]\n",
    "\n",
    "init_coefficients = np.linalg.lstsq(X_pre[:, effective_indices], \n",
    "                                    y_pre, rcond=None)[0].flatten()\n",
    "\n",
    "if com == true_com and mode == 'finetuning': \n",
    "    init_coefficients = avg_weak_coeff\n",
    "    \n",
    "mod, basic_vars = math_eval('+'.join([encoded_feature_names[_] for _ in effective_indices]), \n",
    "                            return_torch=True, split_by_addition=True)\n",
    "init_coefficients, mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "216c1b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# unique to this Burgers' PDE example\n",
    "class ManualICModule(nn.Module):\n",
    "    def __init__(self, *expressions):\n",
    "        super(ManualICModule, self).__init__()\n",
    "        expr1, expr2 = expressions\n",
    "        self.mod0 = sympytorch.SymPyModule(expressions=[expr1])\n",
    "        self.mod1 = sympytorch.SymPyModule(expressions=[expr2])\n",
    "    def forward(self, x_initial):\n",
    "        return self.mod1(x1=self.mod0(x0=x_initial[:, 0]).flatten())\n",
    "\n",
    "class ICModule(nn.Module):\n",
    "    def __init__(self, *expressions):\n",
    "        super(ICModule, self).__init__()\n",
    "        self.mod = sympytorch.SymPyModule(expressions=expressions)\n",
    "    def forward(self, x_initial):\n",
    "        return self.mod(x0=x_initial[:, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "70d17a5a-5c79-4012-aa86-4a44dee55eba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learnable coefficients: Parameter containing:\n",
      "tensor([ 0.1038, -1.0072], 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 = 5 # 5, 10 or 50\n",
    "solver = TorchMLP([2,n_nodes,n_nodes,n_nodes,n_nodes,1], bn=None, \n",
    "                  activation_function=activation_function)\n",
    "# solver = SirenNet(dim_in=2, dim_hidden=50, dim_out=1, num_layers = 4, \n",
    "#                   w0_initial = 30.)\n",
    "\n",
    "physics_calculator = PhysicalConstraintCalculator(symbolic_module=mod, \n",
    "                                                  basic_vars=basic_vars, \n",
    "                                                  init_coefficients=init_coefficients, \n",
    "                                                  learnable_coefficients=True)\n",
    "\n",
    "# ic_module = ICModule(sympify(equation.sympy_format)))\n",
    "# ic_module = ICModule(sympy.exp(recovered_params[0]*((symbols(\"x0\")+recovered_params[1])**2))))\n",
    "ic_module = ManualICModule(symbols(\"x0\")+recovered_params[1], \n",
    "                           sympy.exp(recovered_params[0]*symbols(\"x1\")**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "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, \n",
    "            ic_module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a354caae-01c5-40fb-a672-2d861cb62cf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = True; load = not sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "239324b8-04d0-4c9e-af66-5c77c05876a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.000723435718100518 0.0003336200607009232\n",
      "Epoch 50:  0.00047218045801855624 1.4054084203962702e-05\n",
      "Epoch 100:  0.00047218045801855624 1.4054084203962702e-05\n"
     ]
    }
   ],
   "source": [
    "def closure(return_tuple=False):\n",
    "    if torch.is_grad_enabled():\n",
    "        lbfgs.zero_grad()\n",
    "    l1, l2 = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train, X_train_initial, y_train_initial)\n",
    "    l = torch.add(l1, l2)\n",
    "    if l.requires_grad: \n",
    "        l.backward()\n",
    "    if not return_tuple:\n",
    "        return l\n",
    "    return l1, l2\n",
    "\n",
    "if sim:\n",
    "    flag = False\n",
    "    pinn.set_learnable_ic(flag)\n",
    "    pinn.physics_calculator.set_learnable_coefficients(flag)\n",
    "    lbfgs = torch.optim.LBFGS(pinn.parameters(), \n",
    "                              lr=0.1, max_iter=500, max_eval=500, history_size=300, \n",
    "                              line_search_fn='strong_wolfe')\n",
    "    epochs = 500\n",
    "    best_lt = 1e6; patience = 0\n",
    "    pinn.train()\n",
    "\n",
    "    for i in range(epochs):\n",
    "        lbfgs.step(closure)\n",
    "        \n",
    "        # calculate the loss again for monitoring\n",
    "        if (i%50)==0:\n",
    "            l1, l2 = closure(return_tuple=True)\n",
    "            l1, l2 = l1.item(), l2.item()\n",
    "            lt = l1+l2\n",
    "            if lt < best_lt: best_lt = lt\n",
    "            else: patience += 1\n",
    "            print(\"Epoch {}: \".format(i), l1, l2)\n",
    "            \n",
    "        if patience > 0:\n",
    "            break\n",
    "\n",
    "elif load:\n",
    "    fname = f\"../PMS_data/{denoising_mode}/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/sim{with_initial_data}_pinn{n_nodes}_noiselv{noise_lv}_{smoother_name}{n_components}.pth\"\n",
    "    pinn.load_state_dict(torch.load(fname))\n",
    "    print(\"loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ab3282b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if mode == 'finetuning' and not load:\n",
    "#     fname = f\"../PMS_data/{denoising_mode}/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/sim{with_initial_data}_pinn{n_nodes}_noiselv{noise_lv}_{smoother_name}{n_components}.pth\"\n",
    "#     torch.save(pinn.state_dict(), fname)\n",
    "#     print(\"save\")\n",
    "# else:\n",
    "#     print(\"not save\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "92b6b97b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#base: 111\n",
      "(-31476.718376049364, -31490.28608784265)\n",
      "(-30501.710371521884, -31268.28608784265)\n",
      "(-30492.92651562524, -31266.28608784265)\n"
     ]
    }
   ],
   "source": [
    "### Indecisive ACS ### -> BIC is better!!! (more regularization)\n",
    "pinn.eval()\n",
    "pred = pinn(X_train[:, 0:1], X_train[:, 1:2]).detach().numpy()\n",
    "base = count_parameters(pinn.solver)\n",
    "# why not including pred to u_t in BIC_AIC calculation???\n",
    "assert com == count_parameters(pinn.physics_calculator, False)\n",
    "print(\"#base:\", base)\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), com))\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), base+count_parameters(pinn.physics_calculator, False)))\n",
    "print(BIC_AIC(pred, y_train.detach().numpy(), base+poly_complexities[count_parameters(pinn.physics_calculator, False)-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c3d93dc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'finetuning'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if mode == 'finetuning':\n",
    "    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])\n",
    "elif mode == 'selection':\n",
    "    validation_indices_x = np.array([i for i in range(len(x)) if i>=len(x)//2+1])\n",
    "    validation_indices_t = np.array([i for i in range(len(t)) if i>=len(t)//2+1])\n",
    "mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "478cfe36",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(-31045.660306208752, -31059.188412747448)\n",
      "(-30072.85039331115, -30837.188412747448)\n",
      "(-30064.086340041802, -30835.188412747448)\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",
    "# why not including pred to u_t in BIC_AIC calculation???\n",
    "print(BIC_AIC(val_pred, y_val, com))\n",
    "print(BIC_AIC(val_pred, y_val, base+count_parameters(pinn.physics_calculator, False)))\n",
    "print(BIC_AIC(val_pred, y_val, base+poly_complexities[count_parameters(pinn.physics_calculator, False)-1])) # a good choice with AIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "769bae99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# full_domain_pred = pinn(torch.tensor(X.flatten()[:,None]).float(), \n",
    "#                         torch.tensor(T.flatten()[:,None]).float()).detach().numpy()\n",
    "# full_domain_pred = full_domain_pred.reshape(len(t), len(x)).T\n",
    "\n",
    "# plot_var = full_domain_pred\n",
    "# levels = np.linspace(plot_var.min(), plot_var.max(), 10)\n",
    "# plt.contourf(X, T, plot_var.T, levels)\n",
    "# plt.colorbar(ticks=np.round(levels, 2))\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "bc0dc25f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fname2wsindy = f\"../PMS_data/{denoising_mode}/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/burgers_sim{with_initial_data}_pinn{n_nodes}_noiselv{noise_lv}_{smoother_name}{n_components}_prediction.h5\"\n",
    "# h5file(fname2wsindy, \n",
    "#       {\"usol\": full_domain_pred, \n",
    "#        \"init_coefficients\": init_coefficients\n",
    "#       }, mode='w')\n",
    "# fname2wsindy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "96df7840",
   "metadata": {},
   "outputs": [],
   "source": [
    "# V2 smoother_name='none' (function corrected by LM (levenberg_marquardt))\n",
    "# 2\n",
    "# (6960.518135296146, 6947.005715112373)\n",
    "# (7932.457455495574, 7169.005715112373)\n",
    "# (7941.213665587461, 7171.005715112373)\n",
    "# 3\n",
    "# (7629.728512717052, 7609.4598824413915)\n",
    "# (8601.667832916479, 7831.4598824413915)\n",
    "# (8619.180253100254, 7835.4598824413915)\n",
    "\n",
    "# V2 smoother_name='none' (function corrected by jaxfit with p0 initialized by pysr_params)\n",
    "# 2\n",
    "# (2666.4925604249074, 2652.980140241134)\n",
    "# (3638.4318806243355, 2874.980140241134)\n",
    "# (3647.188090716222, 2876.980140241134)\n",
    "# 3\n",
    "# (13052.25777974368, 13031.98914946802)\n",
    "# (14024.197099943107, 13253.98914946802)\n",
    "# (14041.709520126882, 13257.98914946802)\n",
    "\n",
    "# V2 smoother_name='none' (function corrected by jaxfit)\n",
    "# 2\n",
    "# (4688.597570933143, 4675.08515074937)\n",
    "# (5660.536891132571, 4897.08515074937)\n",
    "# (5669.293101224458, 4899.08515074937)\n",
    "# 3\n",
    "# (9800.22336993561, 9779.954739659948)\n",
    "# (10772.162690135036, 10001.954739659948)\n",
    "# (10789.67511031881, 10005.954739659948)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "cce7fd98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# V1 smoother_name='none' (function directly from pysr)\n",
    "# 2\n",
    "# (5324.327447862634, 5310.81502767886)\n",
    "# (6296.266768062062, 5532.81502767886)\n",
    "# (6305.022978153948, 5534.81502767886)\n",
    "# 3\n",
    "# (3881.683611229945, 3861.414980954285)\n",
    "# (4853.622931429373, 4083.414980954285)\n",
    "# (4871.135351613147, 4087.414980954285)\n",
    "\n",
    "# V1 smoother_name='none' (function corrected by scipy)\n",
    "# 2\n",
    "# (12619.616256354051, 12606.103836170278)\n",
    "# (13591.55557655348, 12828.103836170278)\n",
    "# (13600.311786645365, 12830.103836170278)\n",
    "# 3\n",
    "# (5090.700232076152, 5070.431601800492)\n",
    "# (6062.63955227558, 5292.431601800492)\n",
    "# (6080.151972459354, 5296.431601800492)\n",
    "\n",
    "# V2 smoother_name='none' (function directly from pysr same as if use pysr_params)\n",
    "# 2\n",
    "# (905.3833566363319, 891.8709364525585)\n",
    "# (1877.32267683576, 1113.8709364525585)\n",
    "# (1886.0788869276466, 1115.8709364525585)\n",
    "# 3 | brute\n",
    "# (2536.027946437446, 2515.759316161786)\n",
    "# (3507.9672666368742, 2737.759316161786)\n",
    "# (3525.4796868206477, 2741.759316161786)\n",
    "# 3 | found\n",
    "# (19067.962344248157, 19047.693713972498)\n",
    "# (20039.901664447585, 19269.693713972498)\n",
    "# (20057.41408463136, 19273.693713972498)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "bd477dc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notes: selection | ksvd ดีกว่า dictionary_learning ???\n",
    "\n",
    "# none (no filter) on denoised\n",
    "# 1\n",
    "# (6387.133238243247, 6380.377028151361)\n",
    "# (7359.0725584426755, 6602.377028151361)\n",
    "# (7367.828768534562, 6604.377028151361)\n",
    "# 2\n",
    "# (1538.879435189574, 1525.3670150058006)\n",
    "# (2510.818755389002, 1747.3670150058006)\n",
    "# (2519.5749654808887, 1749.3670150058006)\n",
    "# 3\n",
    "# (2827.1353854996823, 2806.866755224022)\n",
    "# (3799.0747056991104, 3028.866755224022)\n",
    "# (3816.587125882884, 3032.866755224022)\n",
    "# none (no filter) on undenoised\n",
    "# 1\n",
    "# (14322.732737624472, 14315.976527532584)\n",
    "# (15294.672057823898, 14537.976527532584)\n",
    "# (15303.428267915786, 14539.976527532584)\n",
    "# 2\n",
    "# (4232.994498889007, 4219.482078705234)\n",
    "# (5204.9338190884355, 4441.482078705234)\n",
    "# (5213.690029180322, 4443.482078705234)\n",
    "# 3\n",
    "# (8159.8920257322525, 8139.623395456592)\n",
    "# (9131.83134593168, 8361.623395456592)\n",
    "# (9149.343766115453, 8365.623395456592)\n",
    "\n",
    "# lowess on denoised\n",
    "# 1\n",
    "# (12556.936314045266, 12550.180103953378)\n",
    "# (13528.875634244692, 12772.180103953378)\n",
    "# (13537.63184433658, 12774.180103953378)\n",
    "# 2\n",
    "# (7543.226674890819, 7529.714254707045)\n",
    "# (8515.165995090247, 7751.714254707045)\n",
    "# (8523.922205182133, 7753.714254707045)\n",
    "# 3\n",
    "# (10053.498820730207, 10033.230190454546)\n",
    "# (11025.438140929633, 10255.230190454546)\n",
    "# (11051.706771205294, 10261.230190454546)\n",
    "# lowess on undenoised\n",
    "# 1\n",
    "# (23659.065057326072, 23652.308847234184)\n",
    "# (24631.0043775255, 23874.308847234184)\n",
    "# (24639.760587617384, 23876.308847234184)\n",
    "# 2\n",
    "# (3647.0511698673836, 3633.53874968361)\n",
    "# (4618.990490066812, 3855.53874968361)\n",
    "# (4627.746700158698, 3857.53874968361)\n",
    "# 3\n",
    "# (15541.19029720406, 15520.9216669284)\n",
    "# (16513.129617403487, 15742.9216669284)\n",
    "# (16539.39824767915, 15748.9216669284)\n",
    "\n",
    "# kalman on denoised\n",
    "# 1\n",
    "# (13204.5688140811, 13197.812603989212)\n",
    "# (14176.508134280526, 13419.812603989212)\n",
    "# (14185.264344372414, 13421.812603989212)\n",
    "# 2\n",
    "# (3996.366177616859, 3982.853757433086)\n",
    "# (4968.305497816287, 4204.853757433086)\n",
    "# (4977.061707908174, 4206.853757433086)\n",
    "# 3\n",
    "# (10076.267941807173, 10055.999311531512)\n",
    "# (11048.2072620066, 10277.999311531512)\n",
    "# (11074.47589228226, 10283.999311531512)\n",
    "# kalman on undenoised\n",
    "# 1\n",
    "# (20318.512290701114, 20311.756080609226)\n",
    "# (21290.45161090054, 20533.756080609226)\n",
    "# (21299.207820992426, 20535.756080609226)\n",
    "# 2\n",
    "# (-2423.809628596653, -2437.3220487804265)\n",
    "# (-1451.870308397225, -2215.3220487804265)\n",
    "# (-1443.1140983053383, -2213.3220487804265)\n",
    "# 3\n",
    "# (-2076.3332954809202, -2096.6019257565804)\n",
    "# (-1104.3939752814922, -1874.6019257565804)\n",
    "# (-1078.125345005832, -1868.6019257565804)\n",
    "\n",
    "# gaussian on denoised\n",
    "# 1\n",
    "# (26297.802316718964, 26291.046106627076)\n",
    "# (27269.741636918392, 26513.046106627076)\n",
    "# (27278.497847010276, 26515.046106627076)\n",
    "# 2\n",
    "# (6333.820018255945, 6320.3075980721715)\n",
    "# (7305.759338455373, 6542.3075980721715)\n",
    "# (7314.51554854726, 6544.3075980721715)\n",
    "# 3\n",
    "# (7609.801080381209, 7589.532450105549)\n",
    "# (8581.740400580637, 7811.532450105549)\n",
    "# (8608.009030856298, 7817.532450105549)\n",
    "# gaussian on undenoised\n",
    "# 1\n",
    "# (14527.116912351757, 14520.36070225987)\n",
    "# (15499.056232551184, 14742.36070225987)\n",
    "# (15507.812442643071, 14744.36070225987)\n",
    "# 2\n",
    "# (3856.9118887723453, 3843.399468588572)\n",
    "# (4828.851208971773, 4065.399468588572)\n",
    "# (4837.60741906366, 4067.399468588572)\n",
    "# 3\n",
    "# (686.5874853368579, 666.3188550611976)\n",
    "# (1658.5268055362858, 888.3188550611976)\n",
    "# (1684.795435811946, 894.3188550611976)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "9005788c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notes: selection | dictionary_learning แย่กว่า ksvd ไม่เอาแล้ววว\n",
    "# kalman on denoised\n",
    "# 1\n",
    "# 2\n",
    "# (3099.29040612073, 3085.7779859369566)\n",
    "# (4071.229726320158, 3307.7779859369566)\n",
    "# (4079.985936412045, 3309.7779859369566)\n",
    "# 3\n",
    "# (1990.9077120652992, 1970.639081789639)\n",
    "# (2962.847032264727, 2192.639081789639)\n",
    "# (2989.1156625403873, 2198.639081789639)\n",
    "# kalman on undenoised\n",
    "# 1\n",
    "# 2\n",
    "# 3\n",
    "\n",
    "# lowess on denoised\n",
    "# 2\n",
    "# (16625.892628188994, 16612.38020800522)\n",
    "# (17597.83194838842, 16834.38020800522)\n",
    "# (17606.588158480306, 16836.38020800522)\n",
    "# 3\n",
    "# (9591.15180126842, 9570.883170992758)\n",
    "# (10563.091121467845, 9792.883170992758)\n",
    "# (10589.359751743506, 9798.883170992758)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "620babdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "if mode == 'selection': raise SystemExit(\"Exit the program.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed5db4c-235e-42ad-9e94-7f8ee1ad0a2c",
   "metadata": {},
   "source": [
    "#### learn (pinn.physics_calculator.set_learnable_coefficients(True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d3415537-5acf-4f6b-9e73-8166a1c73be3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.00047218045801855624 1.4054084203962702e-05\n",
      "Epoch 50:  0.00047218045801855624 1.4054084203962702e-05\n",
      "Epoch 99:  0.00047218045801855624 1.4054084203962702e-05\n",
      "Epoch 0:  0.00047004083171486855 1.403425085300114e-05\n",
      "Epoch 1000:  0.0004695276147685945 1.2473816241254099e-05\n",
      "Epoch 2000:  0.0004693341033998877 1.2159529433120042e-05\n",
      "Epoch 3000:  0.0004692143702413887 1.2034657629556023e-05\n",
      "Epoch 4000:  0.00046913529513403773 1.1962595635850448e-05\n",
      "Epoch 5000:  0.00046907542855478823 1.1907894986507017e-05\n",
      "Epoch 6000:  0.00046902455505914986 1.1859023288707249e-05\n",
      "Epoch 7000:  0.00046897734864614904 1.1814376193797216e-05\n",
      "Epoch 8000:  0.00046893220860511065 1.1772863217629492e-05\n",
      "Epoch 9000:  0.0004688875051215291 1.1733647625078447e-05\n",
      "Epoch 9999:  0.00046884294715709984 1.1697012268996332e-05\n"
     ]
    }
   ],
   "source": [
    "### using lbfgs ###\n",
    "epochs = 100\n",
    "pinn.train()\n",
    "pinn.physics_calculator.set_learnable_coefficients(True)\n",
    "lbfgs2 = torch.optim.LBFGS(pinn.parameters(), \n",
    "                           lr=0.1, max_iter=500, max_eval=500, history_size=300, \n",
    "                           line_search_fn='strong_wolfe')\n",
    "\n",
    "def closure2(return_tuple=False):\n",
    "    if torch.is_grad_enabled(): \n",
    "        lbfgs2.zero_grad()\n",
    "    l1, l2 = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train, X_train_initial, y_train_initial=None)\n",
    "    l = torch.add(l1, l2)\n",
    "    if l.requires_grad: \n",
    "        l.backward()\n",
    "    if not return_tuple:\n",
    "        return l\n",
    "    return l1, l2\n",
    "\n",
    "for i in range(epochs):\n",
    "    lbfgs2.step(closure2)\n",
    "\n",
    "    # calculate the loss again for monitoring\n",
    "    if (i%50)==0 or i==epochs-1:\n",
    "        pinn.eval()\n",
    "        l1, l2 = closure2(return_tuple=True)\n",
    "        print(\"Epoch {}: \".format(i), l1.item(), l2.item())\n",
    "        pinn.train()\n",
    "\n",
    "### using non-lbfgs ###\n",
    "epochs = 10000\n",
    "pinn.train()\n",
    "flag = not flag\n",
    "pinn.set_learnable_ic(flag)\n",
    "pinn.physics_calculator.set_learnable_coefficients(flag)\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 = torch.add(*pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train))\n",
    "    l.backward()\n",
    "    optimizer.step()\n",
    "    if (i%1000)==0 or i==epochs-1:\n",
    "        pinn.eval()\n",
    "        l1, l2 = pinn.loss(X_train[:, 0:1], X_train[:, 1:2], y_train)\n",
    "        print(\"Epoch {}: \".format(i), l1.item(), l2.item())\n",
    "        pinn.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5b40b39-5aa5-44d1-88b9-cc47fa93c2c1",
   "metadata": {},
   "source": [
    "#### eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8df2aaca-0473-42a0-9cc9-6a5b809d935a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.10290702432394028, -0.9942610859870911]\n",
      "(1.7404578626155824, 1.1665664613246889)\n"
     ]
    }
   ],
   "source": [
    "print(pinn.physics_calculator.coefficients.detach().numpy().tolist())\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": 44,
   "id": "08b8dc05-575e-4bcd-9fc4-7477b1726953",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_domain_pred = pinn(torch.tensor(X.flatten()[:,None]).float(), \n",
    "                        torch.tensor(T.flatten()[:,None]).float()).detach().numpy()\n",
    "full_domain_pred = full_domain_pred.reshape(len(t), len(x)).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "f505bcaf-6505-4f6a-baad-f2c1aedb2241",
   "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": 46,
   "id": "badb0b0a-de7c-48cf-b955-07da0a6dadb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_mean = np.zeros(X_pre.shape)\n",
    "y_mean = np.zeros(y_pre.shape)\n",
    "n_times = 10\n",
    "final_coeffs = np.zeros((n_times, len(effective_indices)))\n",
    "np.random.seed(0)\n",
    "for i 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=weak_pde_lib.spatiotemporal_grid, \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, \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",
    "    final_coeffs[i] = np.linalg.lstsq(X_mean_sub[:, effective_indices], \n",
    "                                      y_mean_sub, rcond=None)[0].flatten()\n",
    "    \n",
    "X_mean = X_mean/n_times\n",
    "y_mean = y_mean/n_times\n",
    "avg_final_coeff = final_coeffs.mean(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c18a370e-cb58-44b2-b7c6-e2a5e45d45e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-50909.726031611426"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# [ 0.10600968 -1.00858733]\n",
    "# (3.4342077507104816, 2.5754747575290824)\n",
    "# -28628.243586182944\n",
    "\n",
    "### ksvd ###\n",
    "# [ 0.1026021  -1.00266824]\n",
    "# (1.4344636753428786, 1.1676400172015544)\n",
    "# -54704.47063053791\n",
    "\n",
    "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",
    "met"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "54813cab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.10280611 -1.00132451]\n",
      "(1.469280012924247, 1.3368286662604492)\n",
      "[ 0.10270856 -1.00099906]\n",
      "(1.404233723878856, 1.3043279945850417)\n"
     ]
    }
   ],
   "source": [
    "##### 0.92 none* (no filter) with learnable lm, ManualIC #####\n",
    "# [ 0.10181832 -1.0003621 ]\n",
    "# (0.9272674937683828, 0.8910571015812302)\n",
    "# [ 0.10184436 -0.99997625]\n",
    "# (0.9233679622411672, 0.9209927171468679)\n",
    "\n",
    "##### 0.92 none* (no filter) with learnable ic (pysr_params), ManualIC #####\n",
    "# [ 0.10280611 -1.00132451]\n",
    "# (1.469280012924247, 1.3368286662604492)\n",
    "# [ 0.10270856 -1.00099906]\n",
    "# (1.404233723878856, 1.3043279945850417)\n",
    "\n",
    "##### none* (no filter) with IC V2 #####\n",
    "# [ 0.10314641 -1.00080487]\n",
    "# (1.6134496391329785, 1.5329627352512405)\n",
    "# [ 0.10306511 -1.00050989]\n",
    "# (1.5580508112907399, 1.5070614218711753)\n",
    "\n",
    "##### none* (no filter) #####\n",
    "### 5 ###\n",
    "# [ 0.10219296 -1.00247914]\n",
    "# (1.2204375307848125, 0.9725237485458504)\n",
    "# [ 0.10220685 -1.00234704]\n",
    "# (1.2207759171724326, 0.986071899533268)\n",
    "\n",
    "##### lowess #####\n",
    "### 5 ###\n",
    "# [ 0.1080395  -0.97994459]\n",
    "# (5.022521331257777, 3.0169806274640374)\n",
    "# [ 0.10796463 -0.97959543]\n",
    "# (5.002545714378356, 2.9620891809463448)\n",
    "\n",
    "##### kalman* > lowess #####\n",
    "### 5 ###\n",
    "# [ 0.10364153 -0.98473121]\n",
    "# (2.5842040484750086, 1.057325444554863)\n",
    "# [ 0.1035684  -0.98446258]\n",
    "# (2.5610710680484745, 1.0073293745517726)\n",
    "### 10 ###\n",
    "# [ 0.10395157 -0.98677679]\n",
    "# (2.636944765181008, 1.3146235419468173)\n",
    "# [ 0.10392218 -0.98648422]\n",
    "# (2.636881098151203, 1.2853034585714322)\n",
    "### 50 ###\n",
    "# [ 0.10402271 -0.9864404 ]\n",
    "# (2.6893361037774426, 1.3333765945868796)\n",
    "# [ 0.10395461 -0.98607194]\n",
    "# (2.673709765076633, 1.2809041887521702)\n",
    "\n",
    "##### gaussian* #####\n",
    "### 5 ###\n",
    "# [ 0.10375497 -0.99599215]\n",
    "# (2.0778778156894164, 1.6770932690532925)\n",
    "# [ 0.10372685 -0.99600383]\n",
    "# (2.063235938549041, 1.6636189818382274)\n",
    "\n",
    "print(mb.params)\n",
    "print(percent_coeff_error(mb.params))\n",
    "print(avg_final_coeff)\n",
    "print(percent_coeff_error(avg_final_coeff))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "a54f6c3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if mode == 'finetuning':\n",
    "#     fname = f\"../PMS_data/ksvd/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/learnable{with_initial_data}_sim_pinn{n_nodes}_noiselv{noise_lv}_{smoother_name}{n_components}.pth\"\n",
    "#     torch.save(pinn.state_dict(), fname)\n",
    "#     print(\"save\", fname)\n",
    "# else:\n",
    "#     print(\"not save\", fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "efdbd21a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAWYAAAD8CAYAAABErA6HAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAcpklEQVR4nO3df7RdZX3n8fcnN0QERNSIP5K0RhvFjCNKY9BaEUFtQEfGqTMTOv6oU5tFl1DsqlODXdXp6h/Wajs6S9pMFqY4S0eWI6gZVwQZrdpOC5OoCAmBehs75BooRhRELOGaz/xx9sWdk/P75z7nfF5r3cXZz/7xPFnc/b3P+T7PfrZsExER1bFs3A2IiIhjJTBHRFRMAnNERMUkMEdEVEwCc0RExSQwR0RUTNvALGmHpHsl7S2VPVHSjZK+Xfz3CcNtZkTEeDWKhXX7Jem/SpqXdKuks0r7Nkm6s9i3tV1dnfSYrwY21ZVtBb5kex3wpWI7ImKaXc3xsbDsAmBd8bMF+AsASXPAlcX+9cDFkta3qqhtYLb9NeC+uuKLgI8Vnz8G/Ot214mImGRNYmHZRcB/d81NwGmSngZsBOZtH7B9BLimOLap5T228Sm27y4ae7ek05sdKGkLtb8ezLH8F09enqzHoPgxJ4yknqOPGc1QxE9H889pa+6R5vuWPXy06T493OLE6NoDi987bPvJ/VzjnHNP9A/ua/7/bMne2x65wXar3nAnVgEHS9sLRVmj8rNbXajXwNwx29uB7QCPP+F0/9ITf3XYVc6MxWetGkk9P1r72KHX8eCq6oxDn/Ld5jfy477zk6b7lv/Dd4fRnJl1/b3b/l+/1/jBfUe5btfKtsc9e83dZ0jaUyraXsSubqhBmVuUN9VrYP4nSU8restPA+7t8ToREVVw2PaGPq+xAKwpba8GDgErmpQ31Ws3ZSfwluLzW4DP9XidiIhpsRN4czE748XA/UXKdzewTtJaSSuAzcWxTbXtMUv6JHAusFLSAvBe4I+BT0n6DeAu4N/286+JmHRJY0y/JrHwBADb24BdwIXAPPAQ8NZi36KkS4EbgDlgh+19repqG5htX9xk1/md/GMi2qlSfrmVVvnlmH4tYuHSfgNvb7JvF7XA3ZHJuCPiONM08BcRx0pgjihpNSMjYlQSmCMiKiaBOcZqUvLLEaOUuyKiT5mREYOWwBzRgczIiFFKYI6xSRojorHcGdFUpspFjEcCc0RExSQwRxR6mcOcgb8YhgTmGIvklyOay90R0UZmZMSoJTBHRFRMAvMEGtUCRsOSNEZEa7lDoqFMlWsvA38xLAnMEREdkLRJ0p2S5iVtbbD/8ZL+l6RvSdon6a2lfZdL2luUv6NdXQnMEWS5z2hN0hxwJXABsB64WNL6usPeDtxu+0xqbzr5U0krJD0P+E1gI3Am8FpJ61rVl8AcIzVp+eXMyIjCRmDe9gHbR4BrgIvqjjHwOEkCTgHuAxaB5wI32X7I9iLwVeD1rSqbrLskImI4VkraU/rZUrd/FXCwtL1QlJV9hFoQPgTcBlxu+yiwFzhH0pMknUTtvYBraKHtO/8i4ngZ+JsM9x89kV0P1mccGrn7sO0NLQ5QgzLXbf8KcAtwHvAs4EZJf217v6T3AzcCDwLfotaTbio95oiI9hY4tpe7mlrPuOytwHWumQe+A5wBYPujts+yfQ61FMe3W1WWwBzHGdZUuUnLL0eU7AbWSVoraQWwGdhZd8xdwPkAkp4CPAc4UGyfXvz354B/A3yyVWVJZUQ0kYG/WGJ7UdKlwA3AHLDD9j5JlxT7twF/BFwt6TZqqY932T5cXOJaSU8CHgHebvsHrepLYI6Zl6ly0Qnbu4BddWXbSp8PAa9ucu7Luqkr3y1jJKYpjZGBvxi26blbIiKmRAJzRETFJDBPmElfWS4i2ktgjmMMY6rcJOaXMyMjxmny7piIAep2RkYG/mIUEpgjIiomgTkiomISmGOoJjG/HDFuuWsi6mTgL8YtgTmiQxn4i1HpKzBL+p3iHVZ7JX1S0omDaliMXl7AGlENPQdmSauA3wY22H4etRWXNg+qYTH5qp5fzuJFUVX93jnLgcdKWg6cxPELR0dERJd6Dsy2vwt8kNri0HcD99v+Yv1xkrYsvUfryNEMqkREtNNPKuMJ1N4SuxZ4OnCypDfWH2d7u+0NtjesWJYc5qwYVRrjodVHH/0ZhGYzMjLwF5I2SbpT0rykrQ32/ydJtxQ/eyX9VNITi31djcf1c/e8EviO7e/ZfgS4DvilPq4X0ZVBBeOIdiTNAVcCFwDrgYslHfOWV9sfsP0C2y8ArgC+avu+Xsbj+gnMdwEvlnSSJFF719X+Pq4XbQxzZblJm5HRKCgnUMcQbQTmbR+wfQS4hlrGoJmLOfa9fl2Nx/WTY74Z+DTwDeC24lrbe71eRKcGEYAzIyPqrFwaCyt+ttTtXwUcLG0vFGXHkXQSsAm4Fjofjyvr651/tt8LvLefa8T0qfo0uZgdDyyeyJcPn9HBkV86bHtDiwPUoMxNjv1XwP+xfR8cNx73Q+B/Snqj7Y83qyx3UEyUYaYrMvAXLSwAa0rbq2mejtjMsWmMrsfjEphjYiSHHGO0G1gnaa2kFdSC7876gyQ9Hng58LlScdfjcX2lMiIiZoHtRUmXAjdQm1Wxw/Y+SZcU+7cVh74e+KLtH5fOvVnS0njcIvBN2ozHJTDHQA0rv9xpb/mh1Uc5aSFfBGPwbO8CdtWVbavbvhq4usG5XY3H5Tc4Jm6qXL+6mZGR/HKMQwJzVN4ocstZgzmqJIE5IqJiEphjYIaRX85MjJhFCcwRERWTwByVNe7ecgb+YlwSmGdcZmRk4C+qJ4E5BmLQ+eVee8uZwxzTIL/FE2KYS35GRLUkMEdEVEwCc1TOuAf9IAN/MV4JzNG3SV5/OQN/UUWTe0dF36o4I2OYveW8tSQmRQJzRETFJDBHZVQhtwzJL8f4JTBHXyY5vxxRVbmrYmZl4C+qKoE5KqEqaYyIZiRtknSnpHlJW5scc66kWyTtk/TVUvlpkj4t6Q5J+yW9pFVdebXUjKrijIxhyoyM6IekOeBK4FXU3pi9W9JO27eXjjkN+HNgk+27JJ1eusSHgettv6F4metJrepLYI6eDSq/XKXecgb+pss/P3IC++95yiAutRGYt30AQNI1wEXA7aVjfg24zvZdALbvLY49FTgH+PWi/AhwpFVlSWVERMBKSXtKP1vq9q8CDpa2F4qysmcDT5D0FUlfl/TmovyZwPeAv5T0TUlXSTq5VWPSY46ZlIG/qHPY9oYW+9WgzHXby4FfBM4HHgv8naSbivKzgMts3yzpw8BW4A+aVZYec4zVINMYWfIzhmgBWFPaXg0canDM9bZ/bPsw8DXgzKJ8wfbNxXGfphaom8pvcvRkkuYvdzrwl/xytLAbWCdpbTF4txnYWXfM54CXSVou6STgbGC/7XuAg5KeUxx3Psfmpo+TVMYEGPRazLM2IyOiX7YXJV0K3ADMATts75N0SbF/m+39kq4HbgWOAlfZ3ltc4jLgE0VQPwC8tVV9CcwxNlWajRHRju1dwK66sm112x8APtDg3FuAVjnsY0zO99GIAcnAX1RdAnOMRdV6y8kvR5UkMEfXpnHgL6JKJucOi4iYEX0F5m4X5ojxq8KMjHGmMZJfjknQ76yMrhbmiBiWPFwS06TnwNzLwhwx+SYpv9ypDPxF1fRzl3W0MIekLUsLgxw5mq+Rs65qszEiqqifwLy0MMdf2H4h8GNqC3Mcw/Z22xtsb1ixbPz5zZgdmZERk6qfwNz1whwx28bdW87AX0yKngNzLwtzxHj1OyMj+eWI0eh3VkZXC3NERER7fQXmbhfmiNk17jRGxCSZvu+mU2bQS35Oo0ZzmOsH/pJfjkmSwBwdSX45Zp2kTZLulDQv6bgZaJLOlXS/pFuKn/eU9v2jpNuK8j3t6sp6zDF0SWPEpJM0B1wJvIrajLTdknbarp/w8Ne2X9vkMq8oXjnV1vR1g6KhKqyRETHBNgLztg8UTzlfA1w0rMrSY46pl/zy7PIjy1g81NESPivrUgzbbW8vba8CDpa2F6i906/eSyR9i9qLWt9pe99SU4AvSjLw3+qufZwE5mirn/zyONIYnTzxl/xy1Dlsu9UMMzUoc932N4Cft/2gpAuBzwLrin0vtX1I0unAjZLusP21ZpUllRER0d4CsKa0vZpar/hRth+w/WDxeRdwgqSVxfah4r/3Ap+hlhppKoE5hmYUveUs9xkjshtYJ2lt8UDdZmBn+QBJT5Wk4vNGavH1+5JOlvS4ovxk4NXAXlpIKiMiog3bi5IuBW4A5oAdtvdJuqTYvw14A/BbkhaBnwCbbVvSU4DPFDF7OfA/bF/fqr4E5mhp0uYvd/JgSfLL0YsiPbGrrmxb6fNHgI80OO8AcGY3dU3WXRc9GcdUucxdjuhdAnNERMUkMMdMSRojJkECczTVa355VGmMdjMy8mBJTKoE5pgaeZVUTIsE5oiIiklgrrBJXIu5yrMxkl+OSZHAPOV6nSo3afOX6yW/HJNssu++qJQq95YjJkkCc0yk+hkZGfiLaZLAHDMh+eWYJAnMcZxe8stVSmMkvxyTLoE5KmP50x8adxMiKiGBeYpNynv+lj/9oUeDci/BuV1+OWmMmDQJzNG3ftIYjQJxu+CcxfFj2uU3PI6R+csRjUnaJOlOSfOStrY47kWSfirpDXXlc5K+Kenz7eqa7LswJlpyyjEpJM0BVwIXAOuBiyWtb3Lc+6m96aTe5cD+TupLYI6+9JrGGFRQTn45RmQjMG/7gO0jwDXARQ2Ouwy4Fri3XChpNfAa4KpOKsurpSJiai070vGYxEpJe0rb221vL22vAg6WtheAs8sXkLQKeD1wHvCiuut/CPg94HGdNCaBOR41qvxyP73lVjdZ8svRh8O2N7TYrwZlrtv+EPAu2z8tXrxaO1F6LXCv7a9LOreTxiQwT6lRTJWr0kMlEUO2AKwpba8GDtUdswG4pgjKK4ELizdmnw28TtKFwInAqZI+bvuNzSpLYI6RGuWAX/LLMUC7gXWS1gLfBTYDv1Y+wPbapc+SrgY+b/uzwGeBK4ryc4F3tgrKkMAcIzTooJyFi2JUbC9KupTabIs5YIftfZIuKfZvG2R9CcwBdJ9frloaI/nlGDbbu4BddWUNA7LtX29S/hXgK+3qynS5GIlB9Ja7eeIvaYyYZH0H5m6eZonOVfm1UlXrLUdMm0H0mDt+miVm0zAG/JJfjmnWV2Du9mmWGI1up8plfYyIaun3jvwQtadZmnZfJG2RtEfSniNHcwNNum7TGONYDyP55Zh0PQfm8tMsrY6zvd32BtsbViybjPWBo3qy1GfMkn5+219K7WmWf6S2oMd5kj4+kFbFVBhWb7mcX04aI6ZRz4HZ9hW2V9t+BrWnYL7c7mmWqJ5u8sujmo2xeOikns9NGiOmQb4fxlBkreWI3g3kyb9On2aJiIj20mOeMsNaVa6bNMage8vlgb9W+eWkMWJaJDDPsEmfvxwxrXJnxkAltxzRvwTmaCtrY0SMVgJzDMwgesv1U+WSX45ZlMA8o5JfjuiOpE2S7pQ0L2lrg/0XSbpV0i3FMhS/XJSfKOn/SvqWpH2S/rBdXVkof4oMY0ZGp2mM5JZjmkmaA64EXkXt/X+7Je20fXvpsC8BO21b0vOBTwFnAA8D59l+UNIJwN9I+oLtm5rVl25TTISkMWLMNgLztg/YPkJtGYqLygfYftD20puzT6Z4i7ZrHizKTyh+6t+wfYz0mKOpcfeWs3BR9GvukY7X7l4paU9pe7vt7aXtVcDB0vYCtbdfH0PS64H3AadTWxJ5qXwO+DrwC8CVtm9u1ZgE5hmU/HLEcQ7b3tBivxqUHdfrtf0Z4DOSzgH+CHhlUf5T4AWSTiv2P8/23maV5Q6Nvgyyt9zL4kVJY8SILABrSturgUPNDrb9NeBZklbWlf+Q2vIVm1pVlsAcDVVp7nKW+YwK2A2sk7RW0gpqK2ruLB8g6Rckqfh8FrAC+L6kJxc9ZSQ9llov+o5WlSWVEZWU/HJUie1FSZcCNwBzwA7b+yRdUuzfBvwq8GZJjwA/Af59MUPjacDHijzzMuBTtlu+vDqBeUp0OlVukPnlcU+RSxojRsn2LmBXXdm20uf3A+9vcN6twAu7qSvdkjhOldIYZUljxKxIYI6eDLq33Gzgr8OpThFTJYG5ghaftWrcTRir5Jdj1uUOmCGd5Jc7SWOMI7ecp/1iliQwR2UljRGzKoE5ujLumRgRsyCBeQoMalW5cc3GKA/8dZJfThojpl0C84wYxPzlcfWWM00uZk0eMImhee5T/+nRz/vveUpX5ya/HLMsPeYABp/GKAflRtu9ShojZkECc3SkmzRGN0G4XX45aYyYRQnMM2CU6y8PomecNEbMugTmaJvG6LS3PKh0RdN2JI0RMyKBecIN4wWs45A0RsTPJDBHS1XpLUeMm6RNku6UNC9pa4P9/0HSrcXP30o6syhfI+mvJO2XtE/S5e3qynS5KdcuvzyI2Ri9BuVGK8o1yy8njRHjVCxyfyXwKmqvmdotaaft20uHfQd4ue0fSLoA2E7tha2LwO/a/oakxwFfl3Rj3bnHSI85+jKsnnLSGFExG4F52wdsHwGuAS4qH2D7b23/oNi8idp7AbF9t+1vFJ9/BOyn9tbtphKYo6lRPemXZT6jAlZK2lP62VK3fxVwsLS9QOvg+hvAF+oLJT2D2ttMbm7VmKQyZli/aYxB95aTxohBW/bw0U6/fR22vaHFfjUoc8MDpVdQC8y/XFd+CnAt8A7bD7RqTAJzNDTs3nKzN5ZA0hhRSQvAmtL2auBQ/UGSng9cBVxg+/ul8hOoBeVP2L6uXWU9f4fsZaQxBqvdVLlhPljSbW+52VoZSWPEhNgNrJO0VtIKYDOws3yApJ8DrgPeZPvvS+UCPgrst/1nnVTWT4+565HGqI5WaYx2veVhDPgljRFVZntR0qXADcAcsMP2PkmXFPu3Ae8BngT8eS0Ws1ikR14KvAm4TdItxSXfXbx1u6GeA7Ptu4G7i88/krQ00pjAHC21SmNEVFURSHfVlW0rfX4b8LYG5/0NjXPUTQ0kx9xqpLEY3dwCcOKyUwZRXYzRIHvLedovorG+E3ztRhptb7e9wfaGFcum4/HhSdAqv9xPGmMYksaIOFZfgbnbkcaYbHnsOmI0+pmV0fVIY1Rbq95yP0G5PCMj0+Qi2uunx7w00niepFuKnwsH1K4YknG9cLXeUn45aYyI4/UzK6PrkcYYnEEv9zms3nKn0luO+JnM7p9Co3xjSTcyTS6iM9W8g2Moeklj9NtbbvTEX9IYEa0lMMdYpsiVJY0RcawE5mgq0+MixiOBeco0yy83S2M06y13GpTPW3kH5628o+1xS/nlpDEi2suyn9GTToJxsxXlypLGiDheeswzrNfecqOg3EmgjojOJDBPoG7nMA/yoZJeAnDSGBHdSWCeIoOYv9yqtzzoXnHSGBGNJTDPqG6nyHUblDvJL0dEYwnMU66bNMYwpscljRHRvQTmGTSM3vKXD5/R1TWTxohJI2mTpDslzUva2mD/GZL+TtLDkt7Zzbn1EpgDaN5b7iWvXJ/GaPfC1fSWo+okzQFXAhcA64GLJa2vO+w+4LeBD/Zw7jESmKdEo4G/RmmMRr3lQQblsvpFi5qlMSImwEZg3vYB20eAa4CLygfYvtf2buCRbs+tlwdMJsygl/tsZpjzkpPGiFHRw490+o1spaQ9pe3ttreXtlcBB0vbC8DZHTaj63MTmGfcoAf8ltIYGfSLCXPY9oYW+xutPe8Or931uUllTKlO0xiNdNtb7mbgL73lmFALwJrS9mrg0LDOTWCeAr0+WNKot5xHqyMa2g2sk7RW0gpgM7BzWOcmlTEjRrHmctIYMa1sL0q6FLgBmAN22N4n6ZJi/zZJTwX2AKcCRyW9A1hv+4FG57aqL4F5CnXyUMmgestJY8SssL0L2FVXtq30+R5qaYqOzm0lqYwZ0ElvedApjHZzlyOiudw9E66X/PIwHr2uT2MsSRojonsJzFOmPo0xzN5y0hgRw5Ec8wQZxMMl9b3lVkH5wlNuf/TzrgebP0Ha7aBfRLSWwDxDOg3K5YAcEaOXVMYU6SWNUa/ToNxPGiP55YjWEpgnWDcDf530lnvpKSeNETF4CcxTqlVveVBBuRMZ9IvoXgLzDGg3Pa7boLyUxuilt5w0RkR7CcxTopxf7qa3nIG+iOpJYJ5y5d5yr0G5PFWuftCvvrdcljRGRG8SmCdUryvKDUqzt2Bn0C+ifwnME6LVwyXN0hiD6C0PUvLLEZ1JYJ4B/QTlRmmMTgb9ksaI6F0C8xRp1lsuy2BfRPX1FZglbZJ0p6R5SVsH1ahZ181X/kZrLzdLYfQTlNv1liNicHq+qyTNAVcCFwDrgYslNV/pJgam04G/foNyq4WL6rVLYyS/HNG5fro7G4F52wdsHwGuAS4aTLOiW8N8dVR6yxGjJbvTN3DXnSi9Adhk+23F9puAs21fWnfcFmBLsfk8YG/vzR2YlcDhtAGoRjuq0AZIO6rWhp+3/eR+LiDpemr/lnYO297UT12D1M+yn2pQdlyUt70d2A4gaY/tDX3UORBVaEcV2lCVdlShDWlH9dowCFUKtt3o57voArCmtL0aONRfcyIiop/AvBtYJ2mtpBXAZmDnYJoVETG7ek5l2F6UdClwAzAH7LC9r81p23utb8Cq0I4qtAGq0Y4qtAHSjrIqtGFm9Tz4FxERw5H5ThERFZPAHBFRMSMPzJJeIOkmSbdI2iNp46jbULTjsuJx8n2S/mQcbSi15Z2SLKmT+ZbDqP8Dku6QdKukz0g6bYR1j/2xfklrJP2VpP3F78Pl42hH0ZY5Sd+U9PkxtuE0SZ8ufif2S3rJuNoyq8bRY/4T4A9tvwB4T7E9UpJeQe0pxefb/hfAB0fdhlJb1gCvAu4aVxuAG4Hn2X4+8PfAFaOotEKP9S8Cv2v7ucCLgbePcXmBy4H9Y6p7yYeB622fAZxZgfbMnHEEZgOnFp8fz3jmPv8W8Me2Hwawfe8Y2rDkvwC/R4OHc0bF9hdtLxabN1Gbkz4KlXis3/bdtr9RfP4RtUC0atTtkLQaeA1w1ajrLrXhVOAc4KMAto/Y/uG42jOrxhGY3wF8QNJBaj3VkfTO6jwbeJmkmyV9VdKLxtAGJL0O+K7tb42j/ib+I/CFEdW1CjhY2l5gDAGxTNIzgBcCN4+h+g9R+yM9ztfAPBP4HvCXRUrlKkknj7E9M6mfR7KbkvS/gac22PX7wPnA79i+VtK/o/aX+ZUjbsNy4AnUvra+CPiUpGd6CHMH27Tj3cCrB11nt+2w/bnimN+n9rX+E6NoEx0+1j8qkk4BrgXeYfuBEdf9WuBe21+XdO4o666zHDgLuMz2zZI+DGwF/mCMbZo5I5/HLOl+4DTbliTgftuntjtvwG24nloq4yvF9j8AL7b9vRG24V8CXwKWloVbeqR9o+17RtWOUnveAlwCnG97eEvVHVvnS4D/bPtXiu0rAGy/bxT117XlBODzwA22/2wM9b8PeBO1P4wnUkv3XWf7jSNux1OBm2w/o9h+GbDV9mtG2Y5ZN45UxiHg5cXn84Bvj6ENny3qRtKzgRWMeCUt27fZPt32M4qbYAE4a0xBeRPwLuB1owrKhUo81l90ED4K7B9HUAawfYXt1cXvwmbgy6MOykU77gEOSnpOUXQ+kNfejNhQUhlt/CbwYUnLgX/mZ0uCjtIOYIekvcAR4C3DSGNMkI8AjwFurMUobrJ9ybAr7fGx/mF4KbXe6m2SbinK3m171xjaUgWXAZ8o/lgeAN465vbMnDySHRFRMXnyLyKiYhKYIyIqJoE5IqJiEpgjIiomgTkiomISmCMiKiaBOSKiYv4/pfLLGW9DR1IAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_var = full_domain_pred\n",
    "levels = np.linspace(plot_var.min(), plot_var.max(), 10)\n",
    "plt.contourf(X, T, plot_var.T, levels)\n",
    "plt.colorbar(ticks=np.round(levels, 2))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "a8441d6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'../PMS_data/ksvd/none/transform_n_nonzero_coefs_none/burgers_learnable_ic_sim_pinn5_noiselv30_none32_prediction.h5'"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fname2wsindy = f\"../PMS_data/ksvd/{smoother_name}/transform_n_nonzero_coefs{transform_n_nonzero}/burgers_learnable{with_initial_data}_sim_pinn{n_nodes}_noiselv{noise_lv}_{smoother_name}{n_components}_prediction.h5\"\n",
    "h5file(fname2wsindy, \n",
    "      {\"usol\": full_domain_pred, \n",
    "       \"avg_final_coeff\": avg_final_coeff\n",
    "      }, mode='w')\n",
    "fname2wsindy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "91f17020-985e-4132-bc49-f17f34470427",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ไม่ใช้แล้ว\n",
    "# noise30 without RDAE\n",
    "# Detect when BIC change is relatively small...\n",
    "# Do not need a high alpha anymore because we regards PINN's interpolation as denoising\n",
    "# 1e-1: -39424.94355880412 |\n",
    "# 1e-2: -40494.82294209987 | (0.9883480862923191, 0.9138094915115704)\n",
    "# 1e-3: -40480.63705933408 | (0.9449528593270584, 0.8633772931360639) ***\n",
    "# 1e-4: -40480.47898107709 | \n",
    "# 1e-5: -40480.47735571946 | \n",
    "# 1e-6: -40480.4773739669 | \n",
    "\n",
    "# noise30 with RDAE\n",
    "# 1e-1: -41925.77786014693 | (1.9317244386014787, 1.6329277859100433)\n",
    "# 1e-2: -44800.92635960881 | (0.6162897470062767, 0.3738864384676466)\n",
    "# 1e-3: -44811.40529297132 | (0.6665730447923304, 0.4183628350028181)\n",
    "# 1e-4: -44811.499217262855 | (0.667096965733599, 0.41882922696416247)\n",
    "# 1e-5: -44811.49996080191 | (0.6671021058580001, 0.41883382329487573)\n",
    "# 1e-6: -44811.49996448646 | (0.6671021470098404, 0.4188338762355748)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e92c22a8-6da9-4470-b9b7-f1580ea9ec59",
   "metadata": {},
   "source": [
    "#### Notes\n",
    "    - 3 main files are required.\n",
    "#### Ideas\n",
    "    - Final ans:  Avg X_pre, y_pre (with 10 random seeds) after PINN training, then OLS\n",
    "    - Final ans from full domain\n",
    "    - Change ps.FiniteDiff (in weak_pde_lib) to Kalman\n",
    "    \n",
    "    - BIC on validation data | full data/domain | calculated after PINN training\n",
    "    - WSINDy as a (better) final ans? | Read the WSINDy paper\n",
    "    - Discover PDE's initial condition -> เสริม DeepONet\n",
    "    - Detect when BIC change is relatively small"
   ]
  }
 ],
 "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
}