Examples/Burgers-PMS-PINN-noise30-IC-lm.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 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 16:03:51,236 [INFO] Remote TPU is not linked into jax; skipping remote TPU.\n",
      "2023-04-17 16:03:51,236 [INFO] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'\n",
      "2023-04-17 16:03:51,237 [INFO] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2023-04-17 16:03:51,238 [INFO] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2023-04-17 16:03:51,239 [INFO] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'\n",
      "2023-04-17 16:03:51,240 [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.00273744  2.01546115]\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 = recovered_params3\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": [
       "'_lm'"
      ]
     },
     "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.0007320275763049722 0.0003197187324985862\n",
      "Epoch 50:  0.0004687905893661082 1.3033405593887437e-05\n",
      "Epoch 100:  0.0004687905893661082 1.3033405593887437e-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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "save\n"
     ]
    }
   ],
   "source": [
    "if mode == 'finetuning' and 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",
      "(-31518.954922194534, -31532.52263398782)\n",
      "(-30543.946917667054, -31310.52263398782)\n",
      "(-30535.16306177041, -31308.52263398782)\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": [
      "(-31019.720872418628, -31033.248978957323)\n",
      "(-30046.910959521025, -30811.248978957323)\n",
      "(-30038.146906251677, -30809.248978957323)\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": "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": 35,
   "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": 36,
   "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": 37,
   "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": 38,
   "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": 40,
   "id": "d3415537-5acf-4f6b-9e73-8166a1c73be3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.00046876369742676616 1.2972490367246792e-05\n",
      "Epoch 50:  0.0004687553446274251 1.2921133020427078e-05\n",
      "Epoch 99:  0.0004687553446274251 1.2921133020427078e-05\n",
      "Epoch 0:  0.0004670634516514838 1.2894578503619414e-05\n",
      "Epoch 1000:  0.0004662036953959614 1.1157544577145018e-05\n",
      "Epoch 2000:  0.0004658316320274025 1.0831438885361422e-05\n",
      "Epoch 3000:  0.0004655779048334807 1.0615878636599518e-05\n",
      "Epoch 4000:  0.00046539847971871495 1.0437272976560052e-05\n",
      "Epoch 5000:  0.00046524452045559883 1.0293438208464067e-05\n",
      "Epoch 6000:  0.00046510473475791514 1.0173831469728611e-05\n",
      "Epoch 7000:  0.0004649859038181603 1.0062473847938236e-05\n",
      "Epoch 8000:  0.00046488599036820233 9.953505468729418e-06\n",
      "Epoch 9000:  0.00046480479068122804 9.842688086791895e-06\n",
      "Epoch 9999:  0.00046473651309497654 9.7311149147572e-06\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": 41,
   "id": "8df2aaca-0473-42a0-9cc9-6a5b809d935a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.10191425681114197, -0.9951547384262085]\n",
      "(1.1993914842605564, 0.7148653268814059)\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": 42,
   "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": 43,
   "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": 44,
   "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": 75,
   "id": "c18a370e-cb58-44b2-b7c6-e2a5e45d45e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-49743.816157590336"
      ]
     },
     "execution_count": 75,
     "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": 76,
   "id": "54813cab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.10181832 -1.0003621 ]\n",
      "(0.9272674937683828, 0.8910571015812302)\n",
      "[ 0.10184436 -0.99997625]\n",
      "(0.9233679622411672, 0.9209927171468679)\n"
     ]
    }
   ],
   "source": [
    "##### 0.92 none* (no filter) with leanable 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": 77,
   "id": "a54f6c3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "save ../PMS_data/ksvd/none/transform_n_nonzero_coefs_none/learnable_lm_sim_pinn5_noiselv30_none32.pth\n"
     ]
    }
   ],
   "source": [
    "if mode == 'finetuning':\n",
    "    fname = f\"../PMS_data/{denoising_mode}/{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": 91,
   "id": "33ba054a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAW8AAAD8CAYAAAC4uSVNAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAdoklEQVR4nO3df7RdZX3n8fcnN4mR32gawCQtkUYwdQBpDFIrIohN0DGltW3SUSljm0mXodJVp4Z2jU5XV5danI7MMm3WXZhCl45ZjIBknCuRYpV2SpgEDJAQItfYkksCIaIgQkku+c4fZ1/cOTm/9z7n7H3O57XWXTn719lfNPub536fZz+PIgIzMyuXaf0OwMzM2ufkbWZWQk7eZmYl5ORtZlZCTt5mZiXk5G1mVkJNk7ekDZIOSNqR2vcaSXdJeiz589TuhmlmZmmttLxvApZW7VsL3B0RC4G7k20zs4FVqyFbdVyS/oekcUkPSbogdWyppN3JsVzyZdPkHRH3AM9U7V4O3Jx8vhn41TyCMTMrsJs4tiGbtgxYmPysAv4GQNIIsC45vghYKWlR1mCmd3jdaRGxHyAi9kuaU+9ESauo/IcwwvRfPH66Kyx5iFfN6Nm9jryqN10jL/fuP6mpkcP1j0176UjdY3qpwYXWlucmnz4YET+T5TsuvmRW/PCZ+v9/Tdnx8OHNEdEoMRMR90g6s8Epy4G/i8pr61sknSLpDOBMYDwi9gBI2pic+0hr/xW1dZq8WxYRo8AowMkz5sQvvebXu33LoTB51tye3evHC17dk/s8P7c4/ecnPFH7gT/x+y/WvWb6956A47sV0fC588D6f836HT985gi3jc1uet4b5u8/R9K21K7RJHe1Yy6wN7U9keyrtf/CNr/7GJ0m76cknZG0us8ADmQNxIZbGRK3DbSDEbE443eoxr5osD+TTp+YTcBVyeergDuyBmJmVnITwPzU9jxgX4P9mbQyVPDLwL3A2ZImJH0Y+DRwuaTHgMuTbTOzYbYJ+FAy6uStwLNJ3+BWYKGkBZJmAiuSczNpWjaJiJV1Dl2W9eZmUKySiVk9SUP2EmC2pAngk8AMgIhYD4wBVwDjwAvA1cmxSUlrgM3ACLAhInZmjafrHZZmg6JpZ6UNtAYN2anjAXykzrExKsk9N27yWEO9GmliZu1x8jZL8UgTKwsnb+sr17vNOuMnx8yshJy8zTJyZ6X1g5N3CfXy1fhuKlPJpNFIE7N+KM/TY2Zmr3DyNkt4pImViZO31eUx3mbF5eRtfVGmencj7qy0fhmMJ8isi9xZaUXk5G1mVkJO3mZmJeTkbT03KPVus37yU2RGZ8ME3Vlp/eTkbTV5mGCFOyttiqSlknZLGpe0tsbxkyX9b0kPStop6erUsY9K2pHsvzaPeJy8radcMrEykjQCrAOWAYuAlZIWVZ32EeCRiDiPyoo7/03STElvAn4PWAKcB7xX0sKsMflJMjNrbgkwHhF7IuIQsBFYXnVOACdKEnAC8AwwCbwR2BIRL0TEJPBt4MqsAXkZNDMbWM8emcXY89UN5Fr2z5a0LbVjNCJGU9tzgb2p7Qngwqov+TyVhYX3AScCvxURRyTtAP5C0muBF6msc7mNjJy8S2ZQZhQsO3dWDpyDEbG4wXHV2BdV278CbAcuBc4C7pL0jxGxS9JngLuA54EHqbTIM3HZxHqmqPVuT0hlLZgA5qe251FpYaddDdwWFePA94FzACLiCxFxQURcTKWc8ljWgIr5NFlfeaRJhUeaWMpWYKGkBZJmAiuolEjSHgcuA5B0GnA2sCfZnpP8+bPArwFfzhqQyyZmZk1ExKSkNcBmYATYEBE7Ja1Ojq8H/hy4SdLDVMosH4+Ig8lX3JrUvA8DH4mIH2aNycnbzKwFETEGjFXtW5/6vA94d51r3553PC6bWE8Utd7dCXdWWhEMzhNlZjZEnLzNanBnpRWdk7cdZdhGmniYoJWVk7d13SDVu82Kwk+VWRvcWWlF4eRdIn413symOHlbV7lkYtYdfrLMqnikiZWBk7e9wiNNzMojU/KW9IfJsj47JH1Z0qy8AjMrGndWWpF0nLwlzQX+AFgcEW+iMlnLirwCs/Jzvduse7I+XdOBV0uaDhzHsfPbmplZF3Q8q2BEPCHps1TmsH0R+EZEfKP6PEmrgFUAs6ad0OntzI7xwryja9bHTWRv6buz0soiS9nkVCoLcC4AXgccL+kD1edFxGhELI6IxTOnDVeHmHVPdeLuNte7TdJSSbsljUtaW+P4f5a0PfnZIellSa9JjuXeP5ilqfIu4PsR8XREHAZuA34pa0DWH3mPNOlHvbvXCd2Gh6QRYB2wDFgErJR01MrGEXF9RJwfEecD1wHfjohnutU/mOUJexx4q6TjkqXuLwN2ZQ3IrJk8krSHCVqblgDjEbEnIg4BG6lUHupZydFLneXeP9hx8o6I+4CvAA8ADyffNZo1ILNG3Lq2LpktaVvqZ1XV8bnA3tT2RLLvGJKOA5YCt0KlfxCY6h/cDzxbq3+wXZmWQYuITwKfzBqENed5TbrPnZWD57nJWXzz4DktnHn3wYhY3OAE1dgXdc7998D/jYhn4Jj+wR8B/0vSByLiiy0EVpcH4lruulXv7ler252VRqWlPT+1PY/6pY8VHF0y6Ur/oJO3mVlzW4GFkhZImkklQW+qPknSycA7gDtSu7vSP+jV460UXOu2foqISUlrgM1URotsiIidklYnx6dWkb+SyjsvP0lde5+kqf7BSeA75NA/6ORtnpDKrAURMQaMVe1bX7V9E3BTjWtz7x902cRy1Y16dy9a3e6stLJx8jZrwJ2VVlRO3lZornWb1ebkbWZWQk7elpu8691udZvV5+Q95DzSpH5npevdVmRO3mZmJeTkXQLDOK+JSyZmjTl5m5mVkJO35SLPzkq3us2ac/K2oebOSisrJ28rlG62uj2niQ0SJ+8hNmzDBM0GiZO3ZdaPxYbNhp2fOiuMonRUut5tZeDkbUPL08BaOyQtlbRb0riktXXOuUTSdkk7JX072Xd2sm/q5zlJ12aNx4sxmJk1IWkEWAdcTmU9y62SNkXEI6lzTgH+GlgaEY9LmgMQEbuB81Pf8wRwe9aY3PK2TPKqd3e7ZOKRJpbREmA8IvZExCFgI5UV4dN+G7gtIh4HiIgDNb7nMuB7EfGvWQNy8h5SHmlidpTZkralflZVHZ8L7E1tTyT70t4AnCrpW5Lul/ShGvepXlm+Yy6bWN8VpaMS3Fk5aP7t8Ax2PXlaK6cejIjFDY6rxr6o2p4O/CKV1vWrgXslbYmI7wIkq86/D7iulYCacfK2jpV5iKA7K61NE8D81PY8YF+Ncw4mK8f/RNI9wHnAd5Pjy4AHIuKpPAIq79NnVuW4Cf91tq7ZCiyUtCBpQa8ANlWdcwfwdknTJR0HXAjsSh1fSU4lE3DLu/AGfTrYXpRM3FlpWUXEpKQ1wGZgBNgQETslrU6Or4+IXZLuBB4CjgA3RsQOgCSZXw78p7xicvI2S7jebY1ExBgwVrVvfdX29cD1Na59AXhtnvH490zrSJnr3WaDwE/gECrKMMF+jTJxZ6UNAidvM7MScvI2w/VuKx8nb2tbHvXuXpVMPNLEBpWTt5lZCWVK3pJOkfQVSY9K2iXporwCM+sGd1baoMg6zvsG4M6IeH/y1tFxOcRkA64bJZMsb1e63m1l1HHylnQScDHwOwDJNImH8gnLuiXrMEGP7zYrhixP4uuBp4G/lfQdSTdKOr76JEmrpqZZPHTEv7IOu16O7XZnpQ2yLMl7OnAB8DcR8WbgJ8AxSwNFxGhELI6IxTOnFePlEBtOrnfbIMmSvCeAiYi4L9n+CpVkblYarndbWXWcvCPiSWCvpLOTXZcBjzS4xEoua727SIsumJVd1tEm1wBfSkaa7AGuzh6SmZk1kyl5R8R2oNHSQWZ94c5KG3Qe9zVE+jmbYDdLJq2M8XZnpQ0aJ29rySCO73ZnpZXZ4D2RA2TQl0AzKxNJSyXtljQu6Zhh0ZIukfSspO3JzydSx/5F0sPJ/m15xONl0KzrPMrEyk7SCLCOyjqUE8BWSZsionqE3T9GxHvrfM07I+JgXjE5edvAqe6sdL17eMXhaUzuy2XKpSXAeETsAZC0EVhOH4dHu2xiTWWpdxe11e16t1WZPTWNR/Kzqur4XGBvansi2VftIkkPSvq6pF9I7Q/gG5Lur/HdHXHL28wMDkZEo2HPqrEvqrYfAH4uIp6XdAXwVWBhcuxtEbFP0hzgLkmPRsQ9WQJ2y3tIFGXR4bxlmQrWrA0TwPzU9jxgX/qEiHguIp5PPo8BMyTNTrb3JX8eAG6nUobJxH/zraFBLJmYdWArsFDSguSN8hXApvQJkk6XpOTzEir59QeSjpd0YrL/eODdwI6sAblsYgOllc5K17utXRExKWkNsBkYATZExE5Jq5Pj64H3A78vaRJ4EVgRESHpNOD2JK9PB/5nRNyZNSYnbzOzFiSlkLGqfetTnz8PfL7GdXuA8/KOx2UT6wqXTMy6y8nb6vIr8WbFNXhPpx1jWEea+OUcG2RO3pa7fpVMPA2sDRMnbzOzEnLytpo6rXcXuaPS9W4bJE7eNpBc77ZB5+RtZlZCTt6Wm16WTKpHmriz0oaNk/eA62SYoMd3mxXf4D2lNvRc77Zh4ORtuSjyKBOzQeTkbaXnercNIydvO0on9e6it7pd77ZB5ORthTD9dS+88tNMozlNXO+2bpG0VNJuSeOS1jY47y2SXpb0/qr9I5K+I+lrecTj5D3AyjAhVasJ26yfJI0A64BlwCJgpaRFdc77DJVFG6p9FNiVV0xO3gU1eVatham7qyglkzyTuUsmlpMlwHhE7ImIQ8BGYHmN864BbgUOpHdKmge8B7gxr4C8ko71TR5J2p2V1si0Qy0vUj1b0rbU9mhEjKa25wJ7U9sTwIXpL5A0F7gSuBR4S9X3fw74Y+DE1iJvzsnb+qIbpRLXuy2DgxGxuMFx1dgXVdufAz4eES8n61VWLpTeCxyIiPslXZIxzlc4eVvHOi2ZZEncLbaizPI2AcxPbc8D9lWdsxjYmCTu2cAVyWLEFwLvk3QFMAs4SdIXI+IDWQJy8jbAr8SbNbEVWChpAfAEsAL47fQJEbFg6rOkm4CvRcRXga8C1yX7LwE+ljVxg5P3wOr2SJN+tLqrud5tvRIRk5LWUBlFMgJsiIidklYnx9c3/IIucPK2geB6t3VbRIwBY1X7aibtiPidOvu/BXwrj3gG73dlK6xejud2ycQGXebknfdbQ9Z77da7OymZ5JG43Vlp9lN5PA25vjVk1grXu23YZUre3XhryAZPt8slrnfbMMra8v4clbeG6jaDJK2StE3StkNH/JD1QjdHmhR9BkFwvduGQ8fJO/3WUKPzImI0IhZHxOKZ04o/UdKw6fb4bk86ZdYdWZ7ct1F5a+hfqEzScqmkL+YSlVmVdGel691mGZJ3RFwXEfMi4kwqbxt9M4+3hqy42i2ZdNrqntx3XMvnVte7XTKxYeGxV2ZmJZTLG5Z5vjVkvdNOvbtXrW4za41b3gOmDKvntKtevdtDBG2YOXlb7vrV6na924aJk7c1VYax3WbDxsnbcuVat1lvOHkPqbIsvtBqvdslExs25XiCrW/aKZnk0epuZ4y3WS9JWippt6RxSWtrHF8u6SFJ25MpQX452T9L0v+T9KCknZL+LI94vBjDABnEkSZmRSBpBFgHXE5lPcutkjZFxCOp0+4GNkVESDoXuAU4B3gJuDQinpc0A/gnSV+PiC1ZYnLLu4Amz5rb7xDa1stat0sm1gdLgPGI2BMRh6hMCbI8fUJEPB8RUyvKH0+yunxUPJ/sn5H8VK883za3vIdQq/XuIo0y8Xwm1omRwy3/3ZktaVtqezQiRlPbc4G9qe0JKqvCH0XSlcCngDlUpsue2j8C3A/8PLAuIu5r+T+iDidvy6xbrW6vnGM9dDAiFjc4rhr7jmk9R8TtwO2SLgb+HHhXsv9l4HxJpyTH3xQRO7IE7KfDzKy5CWB+ansesK/eyRFxD3CWpNlV+39EZSqRpVkDcvK2mlotmeTZ6q430sRDBK0AtgILJS2QNJPKTKqb0idI+nlJSj5fAMwEfiDpZ5IWN5JeTaU1/mjWgFw2GTJlGd9tViQRMSlpDbAZGAE2RMROSauT4+uBXwc+JOkw8CLwW8nIkzOAm5O69zTglojIvGC7k/eA6McwwW6OMHG924omIsaAsap961OfPwN8psZ1DwFvzjsePyF2jCKNMklzycTsp5y8rSO9GtftIYJmtTl5WyGkOytdMjFrzk/JEGmls7KVkkkRZg50ycSGnZO3lYJXzTE7mpP3ABjUCalc7zarz8nbXtGvkkm79W6XTMycvIdGv17OeePpT/HG05/K9B0umZgdy8nbWtZuqzudtNtN4C6ZmDXm5G1dkbW1XY9LJmYVTt4GNK93t9Pq7jRx16p3u2RiVpuT9xDoZb273cRdayZBl0zMmnPyLrleDBNstdXdrVLJK3G4ZGL2Cidvy2UiqqyJ2yUTs/Y4eZuZlZCTtzXUSsmk01Z3O/Vul0zMjubkPeCadVZmLZl0q87tkokVjaSlknZLGpe0tsbx/yDpoeTnnyWdlzq2QdIBSZkWHU5z8ra6ejV7oKeAtaJLljBbBywDFgErJS2qOu37wDsi4lwqK8ePpo7dRA6LDqd5GTTrWN6tbpdMLG/TXjqS129xS4DxiNgDIGkjsBx4ZOqEiPjn1PlbqKwwP3XsHkln5hHIFDd5SizrMMFGJZNmre6sibveSvHgkokV0lxgb2p7ItlXz4eBr3czoI5b3pLmA38HnA4cAUYj4oa8ArPsyrBSvEsmVhCzJW1LbY9GRLrsoRrXRK0vkvROKsn7l3OM7xhZyiaTwB9FxAOSTgTul3RXRDzS7EIrt05a3buePK3hcZdMrM8ORsTiBscngPmp7XnAvuqTJJ0L3Agsi4gf5Bvi0Tpu9kTE/oh4IPn8Y2AXjX+NsALJUjLJyiUTK6GtwEJJCyTNBFYAm9InSPpZ4DbggxHx3W4HlMvvrEkh/s3AfTWOrZK0TdK2Q0f8YJZdnp2ULplYWUTEJLAG2EyloXpLROyUtFrS6uS0TwCvBf5a0vZ0GUbSl4F7gbMlTUj6cNaYMo82kXQCcCtwbUQ8V308qRuNApw8Y07NGpEVR6NWd7fGdLtkYmUQEWPAWNW+9anPvwv8bp1rV+YdT6amj6QZVBL3lyLitnxCsjw06qzMYy6TbnDJxKx1HSdvSQK+AOyKiL/KLyTrl261utOdlVP1bpdMzLLJ8gS9DfggcGlS39ku6Yqc4rImBmXFeJdMzDrTcc07Iv6J2mMfrcDqlUz6UetOc8nErD3+3XUAFfXlnEZDBM2sPcV8yq0wsra6a72cM1XvdsnErHNO3kOk3ZJJL8ol4JKJWSecvK0nXDIxy5eT94Bpt97d61a3SyZm+fB83kMi7xdzLp396FHb3zx4zjHnNJuMClwyMeuUW94llNcY705b3dWJuxmXTMzy5+Q9BPJsdbebuNNcMjHLj5P3AGmn3t1JqztL4q7FJROzzjl5W1dM1bur5zJxq9ssH07eA66dkolb3Wbl4eQ9hGqVTHqZuM3KSNJSSbsljUtaW+P4OZLulfSSpI+1c20nnLytoU4St0smNmgkjQDrgGXAImClpEVVpz0D/AHw2Q6ubZuT94Co1VlZq2TSTqu7Wy1ul0yshJYA4xGxJyIOARuB5ekTIuJARGwFDrd7bSf8kk7JDMo83ma9oJcOt/qb3ez0mpPAaLKE45S5wN7U9gRwYYthZLm2LifvIdKtVnf67UqXTKykDkbE4gbHa61d0OqavFmurctlkwGV5cWcbnZQumRiJTUBzE9tzwP29eDaupy8B0Cniy/0aspXswGwFVgoaYGkmcAKYFMPrq3LZZMh0WiZs7QsrW6XTGxQRcSkpDXAZmAE2BAROyWtTo6vl3Q6sA04CTgi6VpgUUQ8V+varDE5eQ+gVkomtVrdnSTuWrMJ1uOSiZVZRIwBY1X71qc+P0mlJNLStVm5bDIEWml1513nbtbqNrNsnLxLrpN6dzdq3dUlk2ZcMjHLxsm7RFoZ411dMqludedVLgGXTMz6yTXvIdcocV9xwiOvfB57vv7bvO12VJpZdm55D5F2yiXpxJ2VW91m+XPyHiDNSibV6rW6W0nc9UomU63uRlzvNsvOybvE2umsrG51Z0ncafU6Kl0yMesuJ+8B1epLOWl5lkqmuGRi1h1O3kOglVZ3O4l7qmTSSUelSyZm+XDyHhDpenejVnfWxG1mxeDkPeCajTBpJXGnhwm201HpkolZ9zh5l1Qnb1ZWt7qztLg76ah0ycQsP07eA6BeySTd6s4zcZtZ/zl5l0Sey5+1k7hrlUxa6ah0ycSsu5y8B0grrW63uM3yJ+k1ku6S9Fjy56l1zlsqabekcUlrU/t/Q9JOSUckNVqO7RWZkne9QCybdmrDzebuzitxN2t1N+N6tw24tcDdEbEQuDvZPoqkEWAdsAxYBKyUNPWr7Q7g14B7Wr1hx8m7SSDWRc06K/Oa8rXRZFTVXDKxIbccuDn5fDPwqzXOWQKMR8SeiDgEbEyuIyJ2RcTudm6YpeVdNxDrvVpju4vS6jYbAqdFxH6A5M85Nc6ZC+xNbU8k+zqSZUrYWoFcWH2SpFXAqmTzpTsPrN+R4Z55mQ0cLHQMB6q27239ix9L/kyvcHptR3HsT/68u/Wbd6YI/3+A4yhaDD+X9Quem3x6850H1s9u4dRZkraltkcjYjR9gqS/B06vce2fthiOauyLFq89Rpbk3VIgyf8AowCStkVES8X4bipCHEWIoShxFCEGx1G8GPIQEUtz/K531Tsm6SlJZ0TEfklncGzzCyoN3Pmp7XnAvk7jyfJ7b66BmJmV2CbgquTzVcAdNc7ZCiyUtEDSTGAFR/+C3JYsyTvXQMzMSuzTwOWSHgMuT7aR9DpJYwARMQmsATYDu4BbImJnct6VkiaAi4D/I2lzsxt2XDaJiElJU4GMABumAmlgtMnxXilCHEWIAYoRRxFiAMeRVoQYSiMifgBcVmP/PuCK1PYYMFbjvNuB29u5pyI6rpebmVmfeKyXmVkJOXmbmZVQz5O3pPMlbZG0XdI2SUt6HUMSxzXJq/07Jf1lP2JIxfIxSSGplfGo3bj/9ZIelfSQpNslndLDe/d9igVJ8yX9g6Rdyd+Hj/YjjiSWEUnfkfS1PsZwiqSvJH8ndkm6qF+xWH39aHn/JfBnEXE+8Ilku6ckvZPK26DnRsQvAJ/tdQypWOZT6Z1+vF8xAHcBb4qIc4HvAtf14qYFmmJhEvijiHgj8FbgI32c6uGjVEYi9NMNwJ0RcQ5wXgHisRr6kbwDOCn5fDL9GRv++8CnI+IlgIioNaC+V/478MdkeNMqq4j4RjKMCWALlTH7vVCIKRYiYn9EPJB8/jGVZNXxa8udkjQPeA9wY6/vnYrhJOBi4AsAEXEoIn7Ur3isvn4k72uB6yXtpdLi7Ukrr8obgLdLuk/StyW9pQ8xIOl9wBMR8WA/7l/HfwS+3qN75TrXQx4knQm8GbivD7f/HJV/yBtPFdldrweeBv42Kd/cKOn4PsZjdWR5Pb6uJnMAXAb8YUTcKuk3qfwLX/e10y7FMB04lcqvyG8BbpH0+ujCuMkmcfwJ8O6879luHBFxR3LOn1IpIXypFzGR81wPWUk6AbgVuDYinuvxvd8LHIiI+yVd0st7V5kOXABcExH3SbqByvSm/6WPMVkNPR/nLelZ4JSICEkCno2Ik5pdl3MMd1Ipm3wr2f4e8NaIeLqHMfw7KjM+TU0HODW9wJKIeLJXcaTiuQpYDVwWEfWXn8/3nhcB/zUifiXZvg4gIj7Vi/tXxTID+BqwOSL+qg/3/xTwQSr/eM6iUlq8LSI+0OM4Tge2RMSZyfbbgbUR8Z5exmHN9aNssg94R/L5Un46CV4vfTW5N5LeAMykxzOoRcTDETEnIs5MHpQJ4II+Je6lwMeB9/UqcScKMcVC0oj4ArCrH4kbICKui4h5yd+FFcA3e524kzieBPZKOjvZdRng5ZcKqCtlkyZ+D7hB0nTg3/jpdLG9tAHYIGkHcAi4qhslkxL5PPAq4K5KHmNLRKzu9k07nGKhG95GpdX7sKTtyb4/SV5lHkbXAF9K/kHdA1zd53isBr8eb2ZWQn7D0syshJy8zcxKyMnbzKyEnLzNzErIydvMrIScvM3MSsjJ28yshP4/XZLWC/nZGHQAAAAASUVORK5CYII=\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": 92,
   "id": "e82543d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'../PMS_data/ksvd/none/transform_n_nonzero_coefs_none/burgers_learnable_lm_sim_pinn5_noiselv30_none32_prediction.h5'"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fname2wsindy = f\"../PMS_data/{denoising_mode}/{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": 93,
   "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
}