Examples/Burgers1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PDE-FIND for Burger's Equation\n",
    "\n",
    "Samuel Rudy, 2016\n",
    "\n",
    "This notebook demonstrates PDE-FIND on Burger's equation with an added diffusive term.\n",
    "$$\n",
    "u_t + uu_x = 0.1u_{xx}\n",
    "$$\n",
    "The solution given is a single travelling wave, starting out as a Gaussian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%pylab inline\n",
    "pylab.rcParams['figure.figsize'] = (12, 8)\n",
    "import warnings; warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from sympy import symbols\n",
    "\n",
    "import sys; sys.path.append('../')\n",
    "from PDE_FIND import *\n",
    "\n",
    "import sys; sys.path.append('../../parametric-discovery/')\n",
    "from best_subset import *\n",
    "from frols import frols\n",
    "from p_linear_regression import PLinearRegression\n",
    "from r_pca import R_pca\n",
    "from pde_diff_new import RPCA\n",
    "from RobustPCA.rpca import RobustPCA\n",
    "\n",
    "import scipy.io as sio\n",
    "import itertools\n",
    "\n",
    "from sklearn.preprocessing import normalize\n",
    "from sklearn.linear_model import Ridge\n",
    "from sklearn.metrics import mean_squared_error\n",
    "from abess.linear import LinearRegression\n",
    "from bess import PdasLm\n",
    "import pysindy as ps\n",
    "\n",
    "from scipy.signal import savgol_filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = sio.loadmat('../Datasets/burgers.mat')\n",
    "u = real(data['usol'])\n",
    "x = real(data['x'][0])\n",
    "t = real(data['t'][:,0])\n",
    "dt = t[1]-t[0]\n",
    "dx = x[2]-x[1]\n",
    "X, T = np.meshgrid(x, t)\n",
    "XT = np.asarray([X, T]).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STLSQ model: \n",
      "(x0)' = 0.100 x0_11 + -1.001 x0x0_1\n"
     ]
    }
   ],
   "source": [
    "library_functions = [lambda x: x, lambda x: x * x]\n",
    "library_function_names = [lambda x: x, lambda x: x + x]\n",
    "pde_lib = ps.PDELibrary(\n",
    "    library_functions=library_functions,\n",
    "    function_names=library_function_names,\n",
    "    derivative_order=2,\n",
    "    spatial_grid=x,\n",
    "    is_uniform=True,\n",
    ")\n",
    "\n",
    "print('STLSQ model: ')\n",
    "optimizer = ps.STLSQ(threshold=0.1, alpha=1e-5, normalize_columns=True)\n",
    "model = ps.SINDy(feature_library=pde_lib, optimizer=optimizer)\n",
    "model.fit(np.expand_dims(u, -1), t=dt)\n",
    "model.print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct $\\Theta (U)$ and compute $U_t$\n",
    "\n",
    "The function build_linear_system does this for us.  We specify <br>\n",
    "D = highest derivative to appear in $\\Theta$  <br>\n",
    "P = highest degree polynomial of $u$ to appear in $\\Theta$ (not including multiplication by a derivative.  <br>\n",
    "time_diff and space_diff taken via finite differences\n",
    "\n",
    "Printed out is a list of candidate functions for the PDE.  Each is a column of $\\Theta (U)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ut, R, rhs_des = build_linear_system(u, dt, dx, D=3, P=3, time_diff = 'FD', space_diff = 'FD')\n",
    "rhs_des = rhs_des[1:]\n",
    "R = R[:, 1:].real; Ut = Ut.real"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve for $\\xi$\n",
    "\n",
    "TrainSTRidge splits the data up into 80% for training and 20% for validation.  It searches over various tolerances in the STRidge algorithm and finds the one with the best performance on the validation set, including an $\\ell^0$ penalty for $\\xi$ in the loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PDE derived using STRidge\n",
      "u_t = (-1.000403 +0.000000i)uu_{x}\n",
      "    + (0.100145 +0.000000i)u_{xx}\n",
      "   \n"
     ]
    }
   ],
   "source": [
    "# Solve with STRidge\n",
    "w = TrainSTRidge(R,Ut,10**-5,0.1)\n",
    "print(\"PDE derived using STRidge\")\n",
    "print_pde(w, rhs_des)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now identify the same dynamics but with added noise.\n",
    "\n",
    "The only difference from above is that finite differences work poorly for noisy data so here we use polynomial interpolation.  With deg_x or deg_t and width_x or width_t we specify the degree number of points used to fit the polynomials used for differentiating x or t.  Unfortunately, the result can be sensitive to these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# K = 2000 experiments\n",
    "# weak lib only can tolerate up to 75 (maybe more?)\n",
    "# weak lib + KalmanDiff: 90 (maybe more?)\n",
    "# weak lib + KalmanDiff + Robust PCA: 110\n",
    "# set np.random.seed(0) to reproduce...\n",
    "np.random.seed(0)\n",
    "noise_lv = float(120)\n",
    "un = u + 0.01*np.abs(noise_lv)*std(u)*np.random.randn(u.shape[0],u.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "Utn, Rn, rhs_des = build_linear_system(un, dt, dx, D=3, P=3, time_diff = 'poly',\n",
    "                                       deg_x = 4, deg_t = 4, \n",
    "                                       width_x = 10, width_t = 10)\n",
    "rhs_des = rhs_des[1:]\n",
    "Rn = Rn[:, 1:].real; Utn = Utn.real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PDE derived using STRidge\n",
      "u_t = (-0.189497 +0.000000i)uu_{x}\n",
      "   \n"
     ]
    }
   ],
   "source": [
    "# Solve with STRidge\n",
    "w = TrainSTRidge(Rn,Utn,10**-5,1)\n",
    "print(\"PDE derived using STRidge\")\n",
    "print_pde(w, rhs_des)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Goals: implement my own optimizer and differentiation_method class!\n",
    "from functools import partial\n",
    "from tvregdiff import TVRegDiff, tvregdiff, numdiff, pysindydiff, savgol_denoise\n",
    "import sys; sys.path.insert(0, \"../../derivative/\"); import derivative\n",
    "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converged!\n"
     ]
    }
   ],
   "source": [
    "### also works ###\n",
    "# un, _ = RPCA(un, 6e-2)\n",
    "# un, _ = R_pca(un, lmbda=5e-2).fit(max_iter=5000, iter_print=1000)\n",
    "\n",
    "rpca_lambda = None\n",
    "rpca_lambda = 5e-2\n",
    "if rpca_lambda is not None:\n",
    "    rpca = RobustPCA(lamb=rpca_lambda, tol=10, use_fbpca=True, max_iter=int(1e6))\n",
    "    rpca.fit(un)\n",
    "    un = rpca.get_low_rank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# differentiation_method = ps.FiniteDifference\n",
    "# differentiation_kwargs = {}\n",
    "kalpha = 1e-2; poly_deg = 2; diff_order = 2\n",
    "differentiation_method = KalmanDiff\n",
    "differentiation_kwargs = {\"alpha\":kalpha, \"poly_deg\":poly_deg, \"rpca_lambda\":None}\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=XT, \n",
    "                                 include_bias=False, is_uniform=True, K=5000, # new random K points in every calls to the ps.WeakPDELibrary\n",
    "                                 differentiation_method=differentiation_method, \n",
    "                                 differentiation_kwargs=differentiation_kwargs, \n",
    "                                 cache=True\n",
    "                                )\n",
    "\n",
    "# to calculate u, u_1, u_11, ... up to diff_order\n",
    "# weak_pde_lib = ps.WeakPDELibrary(library_functions=[lambda x: x], \n",
    "#                                  function_names=[lambda x: x], \n",
    "#                                  derivative_order=diff_order, p=diff_order, \n",
    "#                                  include_interaction=False, \n",
    "#                                  spatiotemporal_grid=XT, \n",
    "#                                  include_bias=False, is_uniform=True, K=5000, # new random K points in every calls to the ps.WeakPDELibrary\n",
    "#                                  differentiation_method=differentiation_method, \n",
    "#                                  differentiation_kwargs=differentiation_kwargs\n",
    "#                                 )\n",
    "\n",
    "# To have that it's better to calc the coefficients using the weak lib\n",
    "# weak_pde_lib = ps.PDELibrary(library_functions=[lambda x: x, lambda x: x * x], \n",
    "#                         function_names=[lambda x: x, lambda x: x + x], \n",
    "#                         derivative_order=diff_order, spatial_grid=x, \n",
    "#                         include_bias=False, is_uniform=True, \n",
    "#                         differentiation_method=differentiation_method, \n",
    "#                         differentiation_kwargs=differentiation_kwargs\n",
    "#                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "uu = np.expand_dims(un,-1)\n",
    "X_uu_pre1 = weak_pde_lib.fit_transform(uu).reshape(-1, len(weak_pde_lib.get_feature_names()))\n",
    "y_uu_pre1 = weak_pde_lib.convert_u_dot_integral(uu)\n",
    "# for weak_pde_lib = ps.PDELibrary(...)\n",
    "# y_uu_pre1 = weak_pde_lib.differentiation_method(axis=1, **weak_pde_lib.differentiation_kwargs)._differentiate(uu, t).reshape(-1,1)\n",
    "coeffs1 = np.linalg.lstsq(X_uu_pre1[:, [3, 4]], y_uu_pre1, rcond=None)[0]\n",
    "\n",
    "kwargs = {'fit_intercept':False, 'copy_X':True, 'normalize_columns':False}\n",
    "X_uu_pre2, y_uu_pre2, fns = ps_features(uu, t, weak_pde_lib, kwargs)\n",
    "coeffs2 = np.linalg.lstsq(X_uu_pre2[:, [3, 4]], y_uu_pre2, rcond=None)[0]\n",
    "\n",
    "print(np.allclose(X_uu_pre1, X_uu_pre2))\n",
    "print(np.allclose(y_uu_pre1, y_uu_pre2))\n",
    "print(np.allclose(coeffs1, coeffs2))\n",
    "print(np.allclose(weak_pde_lib.cached_xp_full, X_uu_pre1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(u, u^2, u_{x}, u_{xx}, uu_{x}, u^2u_{x}, uu_{xx}, u^2u_{xx})\n",
      "[[ 0.09430853 -0.85176357]]\n",
      "[[ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.09430853]\n",
      " [-0.82664776 -0.85176357]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0.09430853, -0.85176357]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.linear_model import BayesianRidge, LinearRegression as SkLinearRegression\n",
    "from sklearn.base import BaseEstimator\n",
    "from sklearn.pipeline import make_pipeline\n",
    "\n",
    "kwargs = {'fit_intercept':False, 'copy_X':True, 'normalize_columns':False}\n",
    "X_pre, y_pre, fns = ps_features(un, t, weak_pde_lib, kwargs)\n",
    "print(symbols(' '.join(fns)))\n",
    "# reg = BayesianRidge(fit_intercept=False, normalize=True, compute_score=True)\n",
    "reg = SkLinearRegression(fit_intercept=False, normalize=True)\n",
    "print(reg.fit(X_pre[:, [3, 4]], y_pre).coef_)\n",
    "print(bnb(X_pre, y_pre, 2, lam=1e-3, corrected_coefficients=True))\n",
    "\n",
    "class Normalizer(BaseEstimator):\n",
    "    def __init__(self, order=2, axis=0):\n",
    "        super(Normalizer, self).__init__()\n",
    "        self.order = order\n",
    "        self.axis = axis\n",
    "        self.norm = None\n",
    "    def fit(self, X, y):\n",
    "        self.norm = np.divide(1.0, np.linalg.norm(X, ord=self.order, axis=self.axis))\n",
    "    def transform(self, X):\n",
    "        return np.multiply(self.norm, X)\n",
    "    def fit_transform(self, X, y):\n",
    "        self.fit(X, y)\n",
    "        return self.transform(X)\n",
    "\n",
    "# Implement something like NormalizedLinearEstimator...\n",
    "model = make_pipeline(Normalizer(order=2, axis=0), SkLinearRegression(fit_intercept=False))\n",
    "model.fit(X_pre[:, [3, 4]], y_pre)\n",
    "model.steps[1][-1].coef_ = model.steps[0][-1].transform(model.steps[1][-1].coef_)\n",
    "model = model.steps[1][-1]; model.coef_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['u', 'uu', 'u_1', 'u_11', 'uu_1', 'uuu_1', 'uu_11', 'uuu_11']\n",
      "(u)' = 0.094 u_11 + -0.852 uu_1\n",
      "u_t = (0.092865 +0.000000i)u_11\n",
      "    + (-0.844787 +0.000000i)uu_1\n",
      "   \n"
     ]
    }
   ],
   "source": [
    "## give wrong answer\n",
    "# optimizer = ps.STLSQ(threshold=0.1, alpha=1e-5, normalize_columns=True)\n",
    "# optimizer = ps.SR3(threshold=0.1, thresholder='l0', tol=1e-8, normalize_columns=True, max_iter=1000)\n",
    "# optimizer = ABESS(abess_kw={'alpha':1e-2, 'support_size':2, 'cv':5}, group=None, is_normal=False, normalize_columns=True)\n",
    "# optimizer = BESS(bess_kw={'path_type':'seq', 'sequence':[2], 'is_cv':True, 'K':5}, is_normal=True, normalize_columns=False)\n",
    "\n",
    "## give correct answer (for reasonable noise levels) | L0BNB & ps.FROLS are now the best algos. \n",
    "# Fit a weak form model\n",
    "# optimizer = ps.SSR(criteria='model_residual', normalize_columns=True, kappa=1e-2)\n",
    "# optimizer = ps.FROLS(normalize_columns=True, kappa=5e-4, max_iter=100, alpha=1e-1) # kappa = 1e-3 or 5e-4 gives 2 effective candidates\n",
    "optimizer = L0BNB(max_nonzeros=2, lam=1e-3, is_normal=True, normalize_columns=False) # tune lam จาก (X_pre, y_pre) | ตอนนี้ยังไม่ได้ tune\n",
    "# optimizer = BruteForceRegressor(2)\n",
    "# if feature_library=weak_pde_lib, then just differentiation_method=None is fine.\n",
    "n_ensemble_models = 50\n",
    "model = ps.SINDy(feature_library=weak_pde_lib, optimizer=optimizer, \n",
    "                 # differentiation_method=None,\n",
    "                 differentiation_method=KalmanDiff(alpha=kalpha, poly_deg=poly_deg, rpca_lambda=rpca_lambda, axis=0, is_uniform=True), \n",
    "                 cache=True,\n",
    "                 feature_names=['u'])\n",
    "model.fit(np.expand_dims(un, -1), t=dt, ensemble=True, \n",
    "          library_ensemble=True, n_candidates_to_drop=1, n_models=n_ensemble_models)\n",
    "\n",
    "# print the model selected by Pysindy package\n",
    "# (Not giving the best model if set ensemble or library_ensemble = True)\n",
    "print(model.get_feature_names())\n",
    "model.print()\n",
    "\n",
    "# Use the best-performing (on any validation set) coef in model.coef_list (Plse add this functionality to pysindy)\n",
    "# np.einsum('ij,jkl',X_pre,np.array(model.get_coef_list()).T) -> a cool-looking code\n",
    "# from opt_einsum import contract\n",
    "# contract('ij,jkl',X_pre,np.array(model.get_coef_list()).T) -> a cool-looking code\n",
    "print_pde(model.get_coef_list()[np.argmin(np.sum((np.squeeze(np.tensordot(X_pre, np.array(model.get_coef_list()).T, axes=([-1], [0])), axis=1)-y_pre)**2, axis=0))].reshape(-1,1), model.get_feature_names())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "u_t = (0.094309 +0.000000i)u_11\n",
      "    + (-0.851764 +0.000000i)uu_1\n",
      "   \n",
      "u_t = (0.094309 +0.000000i)u_11\n",
      "    + (-0.851764 +0.000000i)uu_1\n",
      "   \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[(<statsmodels.regression.linear_model.RegressionResultsWrapper at 0x7fc7faec6290>,\n",
       "  (3, 4))]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.validate_coef_list()\n",
    "model.validate_coef_list(X_pre, y_pre)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converged!\n",
      "(u)' = 0.074 u_11 + -0.701 uu_1\n"
     ]
    }
   ],
   "source": [
    "kalpha = 0.25\n",
    "pde_lib = ps.PDELibrary(library_functions=library_functions, \n",
    "                        function_names=library_function_names, \n",
    "                        derivative_order=2, spatial_grid=x, \n",
    "                        include_bias=False, is_uniform=True, \n",
    "                        differentiation_method=KalmanDiff, differentiation_kwargs={\"alpha\":kalpha, \"poly_deg\":poly_deg, \"rpca_lambda\":None}\n",
    "                       )\n",
    "optimizer = L0BNB(max_nonzeros=2, lam=1e-3, is_normal=True, normalize_columns=False)\n",
    "# optimizer = BruteForceRegressor(2)\n",
    "model = ps.SINDy(feature_library=pde_lib, optimizer=optimizer, \n",
    "                 # differentiation_method=None, \n",
    "                 differentiation_method=KalmanDiff(alpha=kalpha, poly_deg=poly_deg, rpca_lambda=rpca_lambda, axis=1, is_uniform=True), feature_names=['u'], \n",
    "                )\n",
    "model.fit(np.expand_dims(un, -1), t=dt)\n",
    "model.print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# KalmanDiff -> explicit calc to Theta and u_dot ... (since implementing an extension to ps.BaseDifferentiation is quite confusing for now)\n",
    "# Then, check if they are giving the same results.\n",
    "# Try adding RobustPCA and Savgol (a denoising filter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Further denoise with Savgol?\n",
    "# Combine with WeakPDELibrary? (VERY INTERESTING)\n",
    "# can tolerate up to 120 percent noise!? -> OK, but not really good for low-noise\n",
    "# f = np.expand_dims(un.copy(), -1)\n",
    "f = un.copy()\n",
    "f_t = KalmanDiff(alpha=kalpha, poly_deg=poly_deg, axis=1)._differentiate(f, t).reshape(-1, 1)\n",
    "f_x = KalmanDiff(alpha=kalpha, poly_deg=poly_deg, axis=0)._differentiate(f, x)\n",
    "f_xx = KalmanDiff(alpha=kalpha, poly_deg=poly_deg, axis=0)._differentiate(f_x, x)\n",
    "f_xxx = KalmanDiff(alpha=kalpha, poly_deg=poly_deg, axis=0)._differentiate(f_xx, x)\n",
    "f_xxxx = KalmanDiff(alpha=kalpha, poly_deg=poly_deg, axis=0)._differentiate(f_xxx, x)\n",
    "basis_candidates = [f, f_x, f_xx, f_xxx, f_xxxx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "candidate_lib = []\n",
    "for e in rhs_des:\n",
    "    out = 1\n",
    "    for o, d in get_order_degree(e, accumulate=False):\n",
    "        out = out*(basis_candidates[o]**d)\n",
    "    out = out.ravel().reshape(-1, 1)\n",
    "    candidate_lib.append(out)\n",
    "candidate_lib = np.hstack(candidate_lib)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [-0.20207803  0.        ]\n",
      " [ 0.         -0.70419789]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.07442181]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[[4, 7], [3, 7], [4, 8]]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(bnb(candidate_lib, f_t, 2))\n",
    "# why brute force ไม่เลือกตาม information criterion ที่ใช้ pen also from get_order_degree!!!\n",
    "bf = brute_force(candidate_lib, f_t, 2, top=3, alpha=1.0)\n",
    "[np.nonzero(bf[:, j:j+1])[0].tolist() for j in range(bf.shape[1])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        , -0.20035873,  0.        ],\n",
       "       [-0.67488505,  0.        , -0.6740984 ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.07370067,  0.05917037,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.14215077],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ralpha = 1.0\n",
    "MODEL = Ridge; kwargs = {\"alpha\":ralpha, \"fit_intercept\":False}\n",
    "sk_brute_force(candidate_lib, f_t, 2, MODEL=MODEL, kwargs=kwargs, top=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[-0.70419789],\n",
       "        [ 0.07442181]]),\n",
       " array([25.31060314]),\n",
       " 2,\n",
       " array([20.35498795,  4.80226363]))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.linalg.lstsq(candidate_lib[:, [4, 7]], f_t, rcond=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Applying models to (X_pre, y_pre)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Backward refining selected_inds_list by SelectKBest...\n",
      "SelectKBest\n",
      "(3, 4)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(3, 4, 5)\n",
      "(3, 4)\n",
      "--------------------------------------------------\n",
      "(3, 4, 5, 6)\n",
      "(3, 4, 5)\n",
      "--------------------------------------------------\n",
      "(3, 4, 5, 6, 7)\n",
      "(3, 4, 5, 6)\n",
      "--------------------------------------------------\n",
      "(1, 3, 4, 5, 6, 7)\n",
      "(3, 4, 5, 6, 7)\n",
      "--------------------------------------------------\n",
      "(0, 1, 3, 4, 5, 6, 7)\n",
      "(0, 1, 3, 4, 5, 7)\n",
      "(0, 1, 3, 4, 5)\n",
      "(1, 3, 4, 5)\n",
      "(3, 4, 5)\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{1: ((4,), -30494.901149950994), 2: ((3, 4), -31208.31498852509), 3: ((3, 4, 5), -31317.826723453574), 4: ((3, 4, 5, 6), -31339.434442726182), 5: ((3, 4, 5, 6, 7), -31355.196922923424), 6: ((0, 1, 3, 4, 5, 7), -31374.64243672519), 7: ((0, 1, 3, 4, 5, 6, 7), -31406.615517244514)}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# frols + backward refinement is great!\n",
    "# why (3, 4, 6) -> [4, 6] in noise = 0%\n",
    "coeffs,selected_inds_list,full_coef,_,_ = frols(y_pre,X_pre,max_nonzeros=None)\n",
    "print(\"Backward refining selected_inds_list by SelectKBest...\")\n",
    "backward_refinement(selected_inds_list, (X_pre,y_pre), mode='k best', verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training is over.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[(3, 4, 6),\n",
       " (3, 4),\n",
       " (2, 4),\n",
       " (2, 3, 4, 6, 7),\n",
       " (2, 3, 4, 6),\n",
       " (2, 3, 4, 5, 6, 7),\n",
       " (2, 3, 4),\n",
       " (1, 2, 3, 4, 6),\n",
       " (1, 2, 3, 4, 5, 6, 7),\n",
       " (0, 1, 2, 3, 4, 5, 6, 7)]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pycasso\n",
    "# \"l1\", \"mcp\" and \"scad\"\n",
    "s = pycasso.Solver(normalize(X_pre, axis=0), y_pre.ravel(), \n",
    "                   penalty='scad', gamma=3, lambdas=(100, 1e-2), max_ite=10000)\n",
    "s.train()\n",
    "pycasso_feature_hierarchy = [e for e in sorted(list(set([tuple(np.nonzero(s.coef()['beta'][i])[0]) for i in range(100)]))) if len(e) > 0][::-1]\n",
    "pycasso_feature_hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.          0.          0.          0.         -0.50204916  0.\n",
      "  0.          0.          0.          0.          0.          0.\n",
      " -0.0202092   0.          0.        ]\n",
      "[ 4 12]\n"
     ]
    }
   ],
   "source": [
    "abess_kw={'alpha':None, 'support_size':2, 'cv':5} # Try alphas: 1e2, 1e1, 1e-2, 1e-3, None\n",
    "model = LinearRegression(**abess_kw, covariance_update=True)\n",
    "model.fit(Rn, Utn.ravel(), is_normal=False)\n",
    "print(model.coef_)\n",
    "print(np.where((model.coef_ !=0.0)&(~np.isnan(model.coef_)))[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 1, 2, 3, 4, 6), (0, 2, 4, 5, 7), (1, 2, 3, 4), (1, 3, 6), (3, 6), (3,)]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "abess_feature_hierarchy = []\n",
    "abess_kw = {'alpha':1e2, 'cv':5, 'path_type':'seq'}\n",
    "for ss in range(6, 0, -1): # wide range so that 3 and 4 occur together\n",
    "    model = LinearRegression(support_size=ss, **abess_kw, covariance_update=True)\n",
    "    model.fit(X_pre, y_pre.ravel(), is_normal=True)\n",
    "    # print(model.nonnan_coeffs())\n",
    "    # print(((X_pre@model.nonnan_coeffs()-y_pre.ravel())**2).mean())\n",
    "    abess_feature_hierarchy.append(tuple(np.where((model.nonnan_coeffs() !=0.0)&(~np.isnan(model.nonnan_coeffs())))[0]))\n",
    "abess_feature_hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(0, 1, 3, 4, 5, 6, 7), (4, 5), (4, 5, 6)}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import slir\n",
    "slir_feature_hierarchy = set()\n",
    "for prune_threshold in range(-6, 0):\n",
    "    slr = slir.SparseLinearRegression(n_iter=200, prune_threshold=10**prune_threshold)\n",
    "    slr.fit(X_pre, y_pre)\n",
    "    slir_feature_hierarchy.add(tuple(np.where((slr.coef_!=0.0)&(~np.isnan(slr.coef_)))[0]))\n",
    "slir_feature_hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Backward refinement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Backward refining by SelectKBest...\n",
      "SelectKBest\n",
      "(0, 1, 2, 3, 4, 6)\n",
      "(0, 1, 2, 3, 4)\n",
      "(1, 2, 3, 4)\n",
      "(2, 3, 4)\n",
      "(3, 4)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(0, 2, 4, 5, 7)\n",
      "(2, 4, 5, 7)\n",
      "(4, 5, 7)\n",
      "(4, 7)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(1, 2, 3, 4)\n",
      "--------------------------------------------------\n",
      "(1, 3, 6)\n",
      "(1, 3)\n",
      "(3,)\n",
      "--------------------------------------------------\n",
      "(3, 6)\n",
      "(3,)\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{6: ((0, 1, 2, 3, 4, 6), -31298.21244350962), 5: ((0, 1, 2, 3, 4), -31299.258273874533), 4: ((1, 2, 3, 4), -31289.11105855753), 3: ((2, 3, 4), -31266.028059643904), 2: ((3, 4), -31208.31498852509), 1: ((4,), -30494.901149950994)}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Backward refining by SelectKBest...\")\n",
    "backward_refinement(abess_feature_hierarchy, (X_pre,y_pre), mode='k best', verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RFE refining...\n",
      "RFE_PLinearRegression\n",
      "(0, 1, 2, 3, 4, 6)\n",
      "(0, 1, 2, 3, 4)\n",
      "(1, 2, 3, 4)\n",
      "(2, 3, 4)\n",
      "(3, 4)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(0, 2, 4, 5, 7)\n",
      "(2, 4, 5, 7)\n",
      "(4, 5, 7)\n",
      "(4, 7)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(1, 2, 3, 4)\n",
      "--------------------------------------------------\n",
      "(1, 3, 6)\n",
      "(1, 3)\n",
      "(3,)\n",
      "--------------------------------------------------\n",
      "(3, 6)\n",
      "(3,)\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{6: ((0, 1, 2, 3, 4, 6), -31298.21244350962), 5: ((0, 1, 2, 3, 4), -31299.258273874533), 4: ((1, 2, 3, 4), -31289.11105855753), 3: ((2, 3, 4), -31266.028059643904), 2: ((3, 4), -31208.31498852509), 1: ((4,), -30494.901149950994)}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"RFE refining...\")\n",
    "backward_refinement(abess_feature_hierarchy, (X_pre,y_pre), mode='rfe', verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mRMR refining...\n",
      "mrmr_regression\n",
      "(0, 1, 2, 3, 4, 6)\n",
      "(0, 1, 2, 3, 4)\n",
      "(1, 2, 3, 4)\n",
      "(2, 3, 4)\n",
      "(2, 3)\n",
      "(2,)\n",
      "--------------------------------------------------\n",
      "(0, 2, 4, 5, 7)\n",
      "(2, 4, 5, 7)\n",
      "(4, 5, 7)\n",
      "(4, 7)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(1, 2, 3, 4)\n",
      "--------------------------------------------------\n",
      "(1, 3, 6)\n",
      "(1, 3)\n",
      "(3,)\n",
      "--------------------------------------------------\n",
      "(3, 6)\n",
      "(3,)\n",
      "--------------------------------------------------\n",
      "mrmr_regression\n",
      "(0, 1, 2, 3, 4, 6)\n",
      "(0, 1, 2, 3, 4)\n",
      "(0, 1, 3, 4)\n",
      "(1, 3, 4)\n",
      "(3, 4)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(0, 2, 4, 5, 7)\n",
      "(0, 4, 5, 7)\n",
      "(4, 5, 7)\n",
      "(4, 5)\n",
      "(4,)\n",
      "--------------------------------------------------\n",
      "(1, 2, 3, 4)\n",
      "(2, 3, 4)\n",
      "(3, 4)\n",
      "--------------------------------------------------\n",
      "(1, 3, 6)\n",
      "(1, 3)\n",
      "(3,)\n",
      "--------------------------------------------------\n",
      "(3, 6)\n",
      "(3,)\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{1: ((4,), -30494.901149950994), 2: ((3, 4), -31208.31498852509), 3: ((2, 3, 4), -31266.028059643904), 4: ((1, 2, 3, 4), -31289.11105855753), 5: ((0, 1, 2, 3, 4), -31299.258273874533), 6: ((0, 1, 2, 3, 4, 6), -31298.21244350962)}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"mRMR refining...\")\n",
    "sum([backward_refinement(abess_feature_hierarchy, (X_pre,y_pre), mode='mrmr', sk_normalize_axis=0, verbose=True), \n",
    "    backward_refinement(abess_feature_hierarchy, (X_pre,y_pre), mode='mrmr', sk_normalize_axis=1, verbose=True)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.         -0.07325123]\n",
      " [ 0.          0.09430853  0.0904993 ]\n",
      " [-0.82664776 -0.85176357 -0.61199119]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]]\n",
      "[[ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.09430853]\n",
      " [-0.82664776 -0.85176357]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]]\n",
      "[[ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.09430853]\n",
      " [-0.82664776 -0.85176357]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]\n",
      " [ 0.          0.        ]]\n"
     ]
    }
   ],
   "source": [
    "# lam=1e-1 | 1e-2 | 1e-3 are good choices. -> tune ตาม aic บน (X_pre, y_pre)\n",
    "print(bnb(X_pre, y_pre, 3, lam=1e-1, corrected_coefficients=True))\n",
    "print(bnb(X_pre, y_pre, 3, lam=1e-2, corrected_coefficients=True))\n",
    "print(bnb(X_pre, y_pre, 3, lam=1e-3, corrected_coefficients=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hierarchical selection with BruteForce & ABESS (NOT OPTIMAL: See the include arg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(4,), (2,), (5,)]\n",
      "[(0.00013140094023793248, (4,)), (0.000132715194187583, (2,)), (0.000143364742182267, (5,))]\n",
      "---\n",
      "[(3, 4), (2, 3), (4, 6)]\n",
      "[(0.0001138828733926454, (3, 4)), (0.00011977201108804223, (2, 3)), (0.0001204522748896708, (4, 6))]\n",
      "---\n",
      "[(3, 4, 5), (2, 3, 4), (3, 4, 6)]\n",
      "[(0.00011137113101923466, (3, 4, 5)), (0.00011253090326665838, (2, 3, 4)), (0.00011335473337246112, (3, 4, 6))]\n",
      "---\n",
      "[(3, 4, 5, 6), (1, 3, 4, 5), (0, 3, 4, 5)]\n",
      "[(0.00011084652678552986, (3, 4, 5, 6)), (0.00011085811221648445, (1, 3, 4, 5)), (0.00011107375536617305, (0, 3, 4, 5))]\n",
      "---\n",
      "[(3, 4, 5, 6, 7), (0, 1, 3, 4, 5), (1, 3, 4, 5, 7)]\n",
      "[(0.00011045344356482973, (3, 4, 5, 6, 7)), (0.00011059820909169069, (0, 1, 3, 4, 5)), (0.00011074673771715075, (1, 3, 4, 5, 7))]\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "# training_data = (Rn.copy(), Utn.copy().ravel())\n",
    "# validation_data = (Rn.copy(), Utn.copy().ravel())\n",
    "training_data = (X_pre.copy(), y_pre.copy().ravel())\n",
    "validation_data = (X_pre.copy(), y_pre.copy().ravel())\n",
    "abess_kw = {'cv':5}; n_leaves = 3; nor = True; max_support_sizes = 5\n",
    "\n",
    "prev_nodes = None\n",
    "for ss in range(1, max_support_sizes+1):\n",
    "    if prev_nodes is not None:\n",
    "        ### Brute force ###\n",
    "        effs_collection = set()\n",
    "        for pn in prev_nodes:\n",
    "            model = brute_force(training_data[0], training_data[1], support_size=ss, include=pn, top=n_leaves)\n",
    "            effs = set(tuple(np.nonzero(model[:,j])[0]) for j in range(model.shape[1]))\n",
    "            effs_collection = effs_collection.union(effs)\n",
    "        \n",
    "        ### ABESS ####\n",
    "        # effs_collection = []\n",
    "        # for pn in prev_nodes:\n",
    "        #     selected = frozenset()\n",
    "        #     for nl in range(n_leaves):\n",
    "        #         candidates = Rn.copy()\n",
    "        #         candidates[:, [i for i in range(Rn.shape[1]) if i in selected]] = 0.0\n",
    "        #         if nor: candidates = normalize(candidates, axis=0)\n",
    "        #         model = LinearRegression(support_size=ss, always_select=pn, **abess_kw).fit(candidates, Utn, is_normal=False)\n",
    "        #         effs = frozenset(np.where((model.coef_ != 0.0)&(~np.isnan(model.coef_)))[0])\n",
    "        #         effs_collection.append(effs)\n",
    "        #         effs = effs - frozenset(pn)\n",
    "        #         selected = selected.union(effs)    \n",
    "        # effs_collection = set(effs_collection)\n",
    "        # effs_collection = set(map(tuple, effs_collection))\n",
    "        \n",
    "        prev_nodes = effs_collection\n",
    "\n",
    "    else:\n",
    "        ### Brute force ###\n",
    "        model = brute_force(training_data[0], training_data[1], support_size=ss, top=n_leaves)\n",
    "        prev_nodes = set(tuple(np.nonzero(model[:,j])[0]) for j in range(model.shape[1]))\n",
    "        \n",
    "        ### ABESS ####\n",
    "        # effs_collection = []\n",
    "        # selected = frozenset()\n",
    "        # for nl in range(n_leaves):\n",
    "        #     candidates = Rn.copy()\n",
    "        #     candidates[:, [i for i in range(Rn.shape[1]) if i in selected]] = 0.0\n",
    "        #     if nor: candidates = normalize(candidates, axis=0)\n",
    "        #     model = LinearRegression(support_size=ss, always_select=[], **abess_kw).fit(candidates, Utn)\n",
    "        #     effs = frozenset(np.where((model.coef_ != 0.0)&(~np.isnan(model.coef_)))[0])\n",
    "        #     effs_collection.append(effs)\n",
    "        #     selected = selected.union(effs)\n",
    "        # effs_collection = set(effs_collection)\n",
    "        # prev_nodes = set(map(tuple, effs_collection))\n",
    "        \n",
    "        ### for testing ###\n",
    "        # prev_nodes = {(4,), (3,), (7, )}\n",
    "        \n",
    "    # Performance validation\n",
    "    scores = []\n",
    "    for pn in prev_nodes:\n",
    "        w = np.linalg.lstsq(validation_data[0][:, pn], validation_data[1], rcond=None)[0]\n",
    "        mse = mean_squared_error(validation_data[0][:, pn]@w, validation_data[1])\n",
    "        scores.append((mse, pn))\n",
    "    scores = sorted(scores)\n",
    "    prev_nodes = [v for k,v in scores][:n_leaves]\n",
    "    print(prev_nodes)\n",
    "    print(scores[:n_leaves])\n",
    "    print('---')\n",
    "    prev_nodes = set(prev_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([3, 4, 5, 6, 7]),)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.nonzero(brute_force(X_pre, y_pre, 5).ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Notes\n",
    "    - Hierachical selections helps.\n",
    "    - WeakPDELibrary helps.\n",
    "    - L0BNB helps.\n",
    "    - l0bnb > frols > pycasso > slir > abess\n",
    "    - except l0bnb: best-subset selection methods do not really work with Robust PCA.\n",
    "    - SelectKBest > RFE > mRMR (can be improved?)"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:py3.7]",
   "language": "python",
   "name": "conda-env-py3.7-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.7.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}