KS_exp/KS-PMS-com4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5e664fc2",
   "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\n",
    "\n",
    "import yaml\n",
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b9f61792",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['X_pre', 'best_subsets', 'un', 'y_pre']\n"
     ]
    }
   ],
   "source": [
    "data = sio.loadmat('../Datasets/kuramoto_sivishinky.mat')\n",
    "x = data['x'][:, 0].real\n",
    "t = data['tt'][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",
    "fp1 = \"./IPI_output_files/PMS_data.h5\"\n",
    "fp2 = \"./IPI_output_files/encoded_pde_names.yaml\"\n",
    "X_pre, best_subsets, un, y_pre = h5file(file_path=fp1, mode='r', return_dict=False)\n",
    "encoded_pde_names = OmegaConf.load(fp2)['encoded_pde_names']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e56724bc",
   "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\n",
    "\n",
    "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": 4,
   "id": "2013bef3",
   "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": 5,
   "id": "93c37378",
   "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": 6,
   "id": "1c72d7ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PINN(nn.Module):\n",
    "    def __init__(self, solver, physics_calculator, lb, ub, 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",
    "        # 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",
    "        u_111 = self.gradients(u_11, x)[0]\n",
    "        u_1111 = self.gradients(u_111, x)[0]\n",
    "        physics = self.physics_calculator({nameof(u):u, \n",
    "                                           nameof(u_1):u_1, \n",
    "                                           nameof(u_11):u_11, \n",
    "                                           nameof(u_111):u_111, \n",
    "                                           nameof(u_1111):u_1111})\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 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": 7,
   "id": "0f520a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = default_rng(seed=0)\n",
    "# how about split only by time...\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",
    "domain_dimension = len(sampled_indices_x), len(sampled_indices_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "09febe15",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(MAIN_SEED);\n",
    "torch.manual_seed(MAIN_SEED);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a1c0300e",
   "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": 10,
   "id": "de3e42e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = np.hstack((XX.flatten()[:,None], TT.flatten()[:,None]))\n",
    "y_train = un.T[sampled_indices_t, :][:, sampled_indices_x].flatten()[:,None]\n",
    "lb = torch.tensor(X_train.min(axis=0)).float().requires_grad_(False)\n",
    "ub = torch.tensor(X_train.max(axis=0)).float().requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c4fe67cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'(cos(0.062976345 * x0) * (0.9813761 + sin(0.062976345 * x0)))'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# equation = pd.read_csv(\"./hall_of_fame_2023-05-27_221409.021.csv\").iloc[7][\"Equation\"]\n",
    "\n",
    "equation = \"(cos(0.062976345 * x0) * (0.9813761 + sin(0.062976345 * x0)))\"\n",
    "pysr_params = np.array([0.062976345, 0.062976345, 0.9813761])\n",
    "\n",
    "equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f5dcad2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "func = lambdify(args=sympy.symbols('x0'), expr=equation)\n",
    "# pysr_params = np.array(sorted([float(atom) for atom in sympify(equation).atoms() if atom.is_number]))\n",
    "initial_indices = np.where(X_train[:, 1:2]==0.0)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bb151fc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def initial_function(x, a, b, c): return np.cos(a*x)*(c+np.sin(b*x))\n",
    "def jax_initial_function(x, a, b, c): return jnp.cos(a*x)*(c+jnp.sin(b*x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "39b66bd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-04 15:56:09,886 [INFO] Remote TPU is not linked into jax; skipping remote TPU.\n",
      "2023-06-04 15:56:09,887 [INFO] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'\n",
      "2023-06-04 15:56:09,888 [INFO] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2023-06-04 15:56:09,889 [INFO] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2023-06-04 15:56:09,891 [INFO] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'\n",
      "2023-06-04 15:56:09,892 [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 Parameters ****\n",
      "\n",
      "LM fitting results:\n",
      "----------------------------- \n",
      "parameter      = p1\n",
      "fitted value   = 0.0594\n",
      "standard error = 0.70 %\n",
      "----------------------------- \n",
      "parameter      = p2\n",
      "fitted value   = 0.0592\n",
      "standard error = 0.70 %\n",
      "----------------------------- \n",
      "parameter      = p3\n",
      "fitted value   = 0.9879\n",
      "standard error = 0.04 %\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array([0.06263812, 0.06286727, 0.9839311 ]),\n",
       " array([0.06263812, 0.06286727, 0.9839311 ]),\n",
       " array([0.05936273, 0.05919549, 0.98790786]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm_func = lambda t,p: np.cos(p[0,0]*t)*(p[2,0]+np.sin(p[1,0]*t))\n",
    "\n",
    "recovered_params3 = lm_curve_fit(np.round(pysr_params.reshape(-1, 1), 1), \n",
    "                                 x.flatten(), un[:, 0], \n",
    "                                 lm_func)[0].flatten()\n",
    "\n",
    "# recovered_params1 and recovered_params2 gave similar results. Choose 1.\n",
    "recovered_params2 = np.array(curve_fit(initial_function, x.flatten(), un[:, 0], \n",
    "                                       p0=pysr_params, method='lm')[0])\n",
    "recovered_params1 = np.array(CurveFit().curve_fit(jax_initial_function, x.flatten(), un[:, 0], \n",
    "                                                  p0=pysr_params)[0])\n",
    "\n",
    "recovered_params = recovered_params1\n",
    "recovered_params1, recovered_params2, recovered_params3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b4c56daf",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_function = partial(initial_function, a=recovered_params[0], b=recovered_params[1], c=recovered_params[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "02997ea8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "''"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train_initial, y_train_initial = None, None\n",
    "add_initial_data = 0 # 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": 17,
   "id": "85b0ddd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([64638, 2]), torch.Size([64638, 1]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Converting to tensors\n",
    "X_train = torch.tensor(X_train).float().requires_grad_(True)\n",
    "y_train = torch.tensor(y_train).float().requires_grad_(False)\n",
    "X_train.shape, y_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ebf41777",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-9.95202555e-01,  6.38896462e-04, -9.96653643e-01, -9.96894189e-01]),\n",
       " SymPyModule(expressions=(u_11, u_111, u_1111, u*u_1)),\n",
       " ['u', 'u_1', 'u_11', 'u_111', 'u_1111'])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# com = 3\n",
    "# init_coefficients = np.linalg.lstsq(X_pre[:, np.where(best_subsets[com-1]>0)[0]], \n",
    "#                                     y_pre, rcond=None)[0].flatten()\n",
    "# mod, basic_vars = math_eval(encoded_pde_names[com-1], return_torch=True, split_by_addition=True)\n",
    "\n",
    "com = 4\n",
    "init_coefficients = np.linalg.lstsq(X_pre[:, (3, 4, 5, 6)], \n",
    "                                    y_pre, rcond=None)[0].flatten()\n",
    "mod, basic_vars = math_eval('u_11+u_111+u_1111+u*u_1', return_torch=True, split_by_addition=True)\n",
    "\n",
    "init_coefficients, mod, basic_vars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ce784d91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# unique to this KS PDE example\n",
    "class ManualICModule(nn.Module):\n",
    "    def __init__(self, *expressions):\n",
    "        super(ManualICModule, self).__init__()\n",
    "        raise NotImplementedError\n",
    "    def forward(self, x_initial):\n",
    "        raise NotImplementedError\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": 20,
   "id": "abf3bc1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learnable coefficients: Parameter containing:\n",
      "tensor([-9.9520e-01,  6.3890e-04, -9.9665e-01, -9.9689e-01],\n",
      "       requires_grad=True)\n",
      "[u_11, u_111, u_1111, u*u_1]\n",
      "Basic variables: ['u', 'u_1', 'u_11', 'u_111', 'u_1111']\n"
     ]
    }
   ],
   "source": [
    "activation_function = nn.Tanh()\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",
    "\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",
    "sym_x0 = symbols(\"x0\")\n",
    "expression = sympy.cos(recovered_params[0]*sym_x0)*(recovered_params[2]+sympy.sin(recovered_params[1]*sym_x0))\n",
    "ic_module = ICModule(expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5bf335f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "pinn = PINN(solver, physics_calculator, \n",
    "            lb, ub, \n",
    "            ic_module)\n",
    "# pinn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "fb29ff74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.06804140657186508 0.0026163661386817694\n",
      "Epoch 50:  0.014572379179298878 0.004600754473358393\n",
      "Epoch 100:  0.014572379179298878 0.004600754473358393\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, \n",
    "                       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",
    "\n",
    "ic_flag = False; coeff_flag = True\n",
    "if ic_flag: y_train_initial= None\n",
    "pinn.set_learnable_ic(ic_flag)\n",
    "pinn.physics_calculator.set_learnable_coefficients(coeff_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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "14aeac3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#base: 111\n",
      "(-89851.00843082642, -89887.31466183491)\n",
      "(-88621.51052034086, -89665.31466183491)\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)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "962fdbdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "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])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "db9cc9e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(500212.123834412, 500175.865101102)\n",
      "(501440.3036837652, 500397.865101102)\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)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e3af3a0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (base+count_parameters(pinn.physics_calculator, False)) == count_parameters(pinn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a83ca3c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(pinn.state_dict(), f\"./PMS_PINN_weights/pinn_com{com}.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2e618404",
   "metadata": {},
   "outputs": [],
   "source": [
    "##### add_initial_data = 0 *** #####\n",
    "### base: 111 | com=3 ###\n",
    "## Train\n",
    "# (-92873.31269687883, -92900.5423701352)\n",
    "# (-91643.81478639327, -92678.5423701352)\n",
    "## Val\n",
    "# (337807.93444376386, 337780.74039378134)\n",
    "# (339036.114293117, 338002.74039378134)\n",
    "### base: 111 | com=4 ###\n",
    "## Train\n",
    "# (-89851.00843082642, -89887.31466183491)\n",
    "# (-88621.51052034086, -89665.31466183491)\n",
    "## Val\n",
    "# (500212.123834412, 500175.865101102)\n",
    "# (501440.3036837652, 500397.865101102)\n",
    "\n",
    "##### add_initial_data = 1 (gave wrong selection results) #####\n",
    "### base: 111 | com=3 ###\n",
    "## Train\n",
    "## Val\n",
    "### base: 111 | com=4 ###\n",
    "## Train\n",
    "## Val\n",
    "\n",
    "##### add_initial_data = 2 (gave wrong selection results) #####\n",
    "### base: 111 | com=3 ###\n",
    "## Train\n",
    "## Val\n",
    "### base: 111 | com=4 ###\n",
    "## Train\n",
    "## Val"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}