PDE_discovery/discover_Euler_equ_3cylidner_clean.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from numpy import linalg as LA\n",
    "import pandas as pd\n",
    "import pysindy as ps\n",
    "from tqdm import tqdm\n",
    "from termcolor import colored\n",
    "from scipy.optimize import minimize\n",
    "from sklearn.metrics import r2_score\n",
    "from sklearn.linear_model import LinearRegression\n",
    "import yaml\n",
    "\n",
    "from utils.PolyDiff import PolyDiffPoint\n",
    "from utils.SeqReg import SeqReg\n",
    "\n",
    "%matplotlib inline\n",
    "plt.rcParams[\"font.family\"] = 'Arial'\n",
    "np.set_printoptions(suppress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FitEqu(object):\n",
    "    '''\n",
    "    Sparse regression to fit a equation\n",
    "\n",
    "    1. create dataset\n",
    "    2. preprocess dataset\n",
    "    3. build library\n",
    "    4. fit equation\n",
    "    '''\n",
    "    def __init__(self):\n",
    "        super(FitEqu, self).__init__()\n",
    "\n",
    "    @staticmethod\n",
    "    def load_dataset(case_id, dataset_dir):\n",
    "        '''\n",
    "        load dataset\n",
    "        '''\n",
    "        U = np.load(f'{dataset_dir}/{case_id}/u_3D.npy')\n",
    "        V = np.load(f'{dataset_dir}/{case_id}/v_3D.npy')\n",
    "        P = np.load(f'{dataset_dir}/{case_id}/p_3D.npy')\n",
    "\n",
    "        P = np.transpose(P, axes=[2, 1, 0])\n",
    "        U = np.transpose(U, axes=[2, 1, 0])\n",
    "        V = np.transpose(V, axes=[2, 1, 0])\n",
    "\n",
    "        return U, V, P\n",
    "    \n",
    "    @staticmethod\n",
    "    def select_sub_domain(U, V, P, x_range, y_range):\n",
    "        '''\n",
    "        select a sub-domain\n",
    "        '''\n",
    "        # select a sub-domain\n",
    "        U = U[x_range[0]:x_range[1], y_range[0]:y_range[1], :]\n",
    "        V = V[x_range[0]:x_range[1], y_range[0]:y_range[1], :]\n",
    "        P = P[x_range[0]:x_range[1], y_range[0]:y_range[1], :]            \n",
    "        return U, V, P\n",
    "        \n",
    "    @staticmethod\n",
    "    def add_noise(U, V, P, noise_level=0.01):\n",
    "        Nx, Ny, Nt = U.shape\n",
    "        np.random.seed(0)\n",
    "        U = U + noise_level * np.std(U) * np.random.randn(Nx, Ny, Nt)\n",
    "        V = V + noise_level * np.std(V) * np.random.randn(Nx, Ny, Nt)\n",
    "        P = P + noise_level * np.std(P) * np.random.randn(Nx, Ny, Nt)\n",
    "        return U, V, P\n",
    "        \n",
    "    @staticmethod\n",
    "    def denoise(U, V, P, SVD_mode, is_show=False):\n",
    "        '''\n",
    "        SVD denoise\n",
    "        '''\n",
    "        Nx, Ny, Nt = U.shape\n",
    "        U = U.reshape(Nx * Ny, Nt)\n",
    "        V = V.reshape(Nx * Ny, Nt)\n",
    "        P = P.reshape(Nx * Ny, Nt)\n",
    "\n",
    "        uu, sigmau, vu = np.linalg.svd(U, full_matrices=False); vu = vu.T\n",
    "        uv, sigmav, vv = np.linalg.svd(V, full_matrices=False); vv = vv.T\n",
    "        up, sigmap, vp = np.linalg.svd(P, full_matrices=False); vp = vp.T\n",
    "\n",
    "        # singular values of clean data\n",
    "        if is_show:\n",
    "            fig = plt.figure()\n",
    "            plt.semilogy(np.arange(Nt), sigmau, 'g', label = r'$u$')\n",
    "            plt.semilogy(np.arange(Nt), sigmav, 'b', label = r'$v$')\n",
    "            plt.semilogy(np.arange(Nt), sigmap, 'r', label = r'$p$')\n",
    "\n",
    "            plt.legend(loc = 'upper right', fontsize = 20)\n",
    "            plt.xticks(fontsize = 16)\n",
    "            plt.yticks(fontsize = 16)\n",
    "            plt.xlim([0, 30])\n",
    "\n",
    "        mode_u, mode_v, mode_p = SVD_mode\n",
    "        U = uu[:,0:mode_u].dot(np.diag(sigmau[0:mode_u]).dot(vu[:,0:mode_u].T)).reshape(Nx, Ny, Nt)\n",
    "        V = uv[:,0:mode_v].dot(np.diag(sigmav[0:mode_v]).dot(vv[:,0:mode_v].T)).reshape(Nx, Ny, Nt)\n",
    "        P = up[:,0:mode_p].dot(np.diag(sigmap[0:mode_p]).dot(vp[:,0:mode_p].T)).reshape(Nx, Ny, Nt)\n",
    "\n",
    "        return U, V, P\n",
    "\n",
    "    @staticmethod\n",
    "    def non_dimensionlize(U, V, P, dx, dy, dt, ref_params, sim_params):\n",
    "        '''\n",
    "        non-dimensionlize dx, dy, and dt\n",
    "        '''\n",
    "        # reference parameters\n",
    "        l_ref, v_ref, p_ref = ref_params['l_ref'], ref_params['v_ref'], ref_params['p_ref'] \n",
    "        t_ref, w_ref = l_ref / v_ref, v_ref / l_ref\n",
    "\n",
    "        # non-dimensionlize data\n",
    "        U, V, P = U/v_ref, V/v_ref, P/p_ref\n",
    "        dx, dy, dt = dx / l_ref, dy / l_ref, dt / t_ref\n",
    "\n",
    "        # analyze the dimensionless numbers\n",
    "        Re = v_ref * l_ref * sim_params['rho'] / sim_params['mu']\n",
    "        Fr = ref_params['v_ref'] / float(np.sqrt(sim_params['g'] * ref_params['l_ref']))\n",
    "        Eu = ref_params['p_ref'] / (sim_params['rho'] * ref_params['v_ref']**2)\n",
    "        pi_group = {'Re': Re, 'Fr': Fr, 'Eu': Eu}\n",
    "        print(colored(f'Re: {int(Re)}, coef_best: {round(1 / Re, 6)}', 'red'))\n",
    "        print(colored(f'Eu: {int(Eu)}, coef_best: {round(Eu, 6)}', 'red'))\n",
    "        # print(colored(f'Fr: {int(Fr)}, coef_best: {round(1/Fr/Fr, 6)}', 'red'))\n",
    "        return U, V, P, dx, dy, dt, pi_group\n",
    "\n",
    "    @staticmethod\n",
    "    def sample_points(Nx, Ny, boundary_x, boundary_y, boundary_t):\n",
    "        '''\n",
    "        sample a certain number of points to train\n",
    "        '''\n",
    "        np.random.seed(0)\n",
    "        points, count = {}, 0\n",
    "\n",
    "        num_xy = 500\n",
    "        for _ in range(num_xy):\n",
    "            x = np.random.choice(np.arange(boundary_x, Nx-boundary_x), 1)[0]\n",
    "            y = np.random.choice(np.arange(boundary_y, Ny-boundary_y), 1)[0]\n",
    "            for t in range(10, 100, 10):\n",
    "                points[count] = [x, y, t]\n",
    "                count = count + 1\n",
    "        return points\n",
    "\n",
    "    @staticmethod\n",
    "    def cal_derivatives(U, V, P, points, boundary_x, boundary_y, boundary_t, dx, dy, dt, deg=5):\n",
    "        num_points = len(points)\n",
    "        p = np.zeros((num_points,1))\n",
    "        u = np.zeros((num_points,1))\n",
    "        v = np.zeros((num_points,1))\n",
    "        ut, vt = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "        ux, vx = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "        uy, vy = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "        uxx, vxx = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "        uxy, vxy = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "        uyy, vyy = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "        px, py = np.zeros((num_points,1)), np.zeros((num_points,1))\n",
    "\n",
    "        Nx_sample, Ny_sample, Nt_sample = 2*boundary_x-1, 2*boundary_y-1, 2*boundary_t-1\n",
    "\n",
    "        for idx_p in tqdm(points.keys()):\n",
    "            [x, y, t] = points[idx_p]\n",
    "            u[idx_p] = U[x, y, t]\n",
    "            v[idx_p] = V[x, y, t]\n",
    "            p[idx_p] = P[x, y, t]\n",
    "            \n",
    "            ut_part = U[x, y, t-int((Nt_sample-1)/2):t+int((Nt_sample+1)/2)]\n",
    "            ux_part = U[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y, t]\n",
    "            uy_part = U[x, y-int((Ny_sample-1)/2):y+int((Ny_sample+1)/2), t]\n",
    "            ux_part_yp = U[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y+1, t]\n",
    "            ux_part_ym = U[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y-1, t]\n",
    "\n",
    "            vt_part = V[x, y, t-int((Nt_sample-1)/2):t+int((Nt_sample+1)/2)]\n",
    "            vx_part = V[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y, t]\n",
    "            vy_part = V[x, y-int((Ny_sample-1)/2):y+int((Ny_sample+1)/2), t]\n",
    "            vx_part_yp = V[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y+1, t]\n",
    "            vx_part_ym = V[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y-1, t]\n",
    "\n",
    "            px_part = P[x-int((Nx_sample-1)/2):x+int((Nx_sample+1)/2), y, t]\n",
    "            py_part = P[x, y-int((Ny_sample-1)/2):y+int((Ny_sample+1)/2), t]\n",
    "\n",
    "            ut[idx_p] = PolyDiffPoint(ut_part, np.arange(Nt_sample)*dt, deg, 1)[0]\n",
    "            ux_diff = PolyDiffPoint(ux_part, np.arange(Nx_sample)*dx, deg, 2)\n",
    "            uy_diff = PolyDiffPoint(uy_part, np.arange(Ny_sample)*dy, deg, 2)\n",
    "            ux_diff_yp = PolyDiffPoint(ux_part_yp, np.arange(Nx_sample)*dx, deg, 2)\n",
    "            ux_diff_ym = PolyDiffPoint(ux_part_ym, np.arange(Nx_sample)*dx, deg, 2)\n",
    "\n",
    "            vt[idx_p] = PolyDiffPoint(vt_part, np.arange(Nt_sample)*dt, deg, 1)[0]\n",
    "            vx_diff = PolyDiffPoint(vx_part, np.arange(Nx_sample)*dx, deg, 2)\n",
    "            vy_diff = PolyDiffPoint(vy_part, np.arange(Ny_sample)*dy, deg, 2)\n",
    "            vx_diff_yp = PolyDiffPoint(vx_part_yp, np.arange(Nx_sample)*dx, deg, 2)\n",
    "            vx_diff_ym = PolyDiffPoint(vx_part_ym, np.arange(Nx_sample)*dx, deg, 2)\n",
    "\n",
    "            ux[idx_p], uxx[idx_p] = ux_diff[0], ux_diff[1]\n",
    "            uy[idx_p], uyy[idx_p] = uy_diff[0], uy_diff[1]\n",
    "            uxy[idx_p] = (ux_diff_yp[0] - ux_diff_ym[0]) / (2 * dy)\n",
    "\n",
    "            vx[idx_p], vxx[idx_p] = vx_diff[0], vx_diff[1]\n",
    "            vy[idx_p], vyy[idx_p] = vy_diff[0], vy_diff[1]\n",
    "            vxy[idx_p] = (vx_diff_yp[0] - vx_diff_ym[0]) / (2 * dy)\n",
    "\n",
    "            px[idx_p] = PolyDiffPoint(px_part, np.arange(Nx_sample)*dx, deg, 1)[0]\n",
    "            py[idx_p] = PolyDiffPoint(py_part, np.arange(Ny_sample)*dy, deg, 1)[0]\n",
    "        \n",
    "        base_library = [u, v, p, ut, ux, uy, uxx, uxy, uyy, vt, vx, vy, vxx, vxy, vyy, px, py]\n",
    "        return base_library\n",
    "\n",
    "    @staticmethod\n",
    "    def parse_library(base_library, sim_params, ref_params):\n",
    "        '''\n",
    "        calculate derivatives and build library for sparse regression\n",
    "        '''\n",
    "        u, v, p, ut, ux, uy, uxx, uxy, uyy, vt, vx, vy, vxx, vxy, vyy, px, py = base_library\n",
    "        g_x = np.zeros_like(u)\n",
    "        g_y = np.ones_like(u) * sim_params['g'] * ref_params['l_ref'] / ref_params['v_ref']**2\n",
    "        # ########################standard sindy library#########################\n",
    "        # X_library = [u*ux, v*uy, uxx, uyy, px]\n",
    "        # names = ['u*ux', 'v*uy', 'uxx', 'uyy', 'px']\n",
    "\n",
    "        # #########################rotationial invariance#########################\n",
    "        ########################## More terms (6 terms in total, 4 independent terms)##########################\n",
    "        ut_ = np.concatenate([ut+vt, vt+ut])\n",
    "        u_ = np.concatenate([u+v, v+u])\n",
    "        uu_ = np.concatenate([u*(u+v), v*(v+u)])\n",
    "        vu_ = np.concatenate([v*(u+v), u*(v+u)])\n",
    "\n",
    "        ux_ = np.concatenate([ux+vx, vy+uy])\n",
    "        uy_ = np.concatenate([uy+vy, vx+ux])\n",
    "        uxx_ = np.concatenate([uxx+vxx, vyy+uyy])\n",
    "        uyy_ = np.concatenate([uyy+vyy, vxx+uxx])\n",
    "\n",
    "        uux_ = np.concatenate([u*(ux+vx), v*(vy+uy)])\n",
    "        uuy_ = np.concatenate([u*(uy+vy), v*(vx+ux)])\n",
    "        uuxx_ = np.concatenate([u*(uxx+vxx), v*(vyy+vyy)])\n",
    "        uuyy_ = np.concatenate([u*(uyy+vyy), v*(vxx+uxx)])\n",
    "\n",
    "        vux_ = np.concatenate([v*(ux+vx), u*(vy+uy)])\n",
    "        vuy_ = np.concatenate([v*(uy+vy), u*(vx+ux)])\n",
    "        vuxx_ = np.concatenate([v*(uxx+vxx), u*(vyy+vyy)])\n",
    "        vuyy_ = np.concatenate([v*(uyy+vyy), u*(vxx+uxx)])\n",
    "\n",
    "        uuxx_ = np.concatenate([u*(uxx+vxx), v*(vyy+vyy)])\n",
    "        uuyy_ = np.concatenate([u*(uyy+vyy), v*(vxx+uxx)])\n",
    "        vuxx_ = np.concatenate([v*(uxx+vxx), u*(vyy+vyy)])\n",
    "        vuyy_ = np.concatenate([v*(uyy+vyy), u*(vxx+uxx)])\n",
    "        \n",
    "        p_ = np.concatenate([p, p])\n",
    "        px_ = np.concatenate([px+py, py+px])\n",
    "\n",
    "        ut, u, uu, vu, ux, uy, uxx, uyy = ut_, u_, uu_, vu_, ux_, uy_, uxx_, uyy_\n",
    "        uux, uuy, uuxx, uuyy, vux, vuy, vuxx, vuyy, px = uux_, uuy_, uuxx_, uuyy_, vux_, vuy_, vuxx_, vuyy_, px_\n",
    "        p = p_\n",
    "\n",
    "        # term and position: (uux: 7, vuy: 12, uxx: 5, uyy: 6, px: 16)\n",
    "        X_library = [\n",
    "            u, uu, vu, ux, uy, \n",
    "            uxx, uyy, uux, uuy, uuxx, \n",
    "            uuyy, vux, vuy, vuxx, vuyy, \n",
    "            p, px]\n",
    "        names = ['u', 'uu', 'vu', 'ux', 'uy', 'uxx', 'uyy', 'uux', \n",
    "                'uuy', 'uuxx', 'uuyy', 'vux', 'vuy', 'vuxx', 'vuyy', 'p', 'px']\n",
    "        \n",
    "        ##########################Reshape data##########################\n",
    "        X_library = np.squeeze(np.stack(X_library, axis=-1))\n",
    "        y_library = ut.reshape(-1, 1)\n",
    "\n",
    "        return X_library, y_library, names\n",
    "\n",
    "    @staticmethod\n",
    "    def normalization(X_library, y_library):\n",
    "        '''\n",
    "        Rescale the data by each column\n",
    "        rescale the data by the absolute mean for each column\n",
    "        '''\n",
    "        norm_coef = np.mean(np.abs(np.mean(X_library, axis=0)))\n",
    "        X_library = X_library / norm_coef\n",
    "        y_library = y_library / norm_coef\n",
    "        return X_library, y_library, norm_coef\n",
    "\n",
    "    @staticmethod\n",
    "    def check_library(X_library, y_library):\n",
    "        '''\n",
    "        check whether the library has any Nan\n",
    "        '''\n",
    "        if np.any(np.isnan(X_library)) or np.any(np.isnan(y_library)):\n",
    "            print('Nan exists in library.')\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "    @staticmethod\n",
    "    def check_equ(X_library, y_library, pi_group):\n",
    "        '''\n",
    "        check the r2 for the target equation\n",
    "        '''\n",
    "        pred_best = - X_library[:, 7] - X_library[:, 12]\\\n",
    "                    + 1/pi_group['Re'] * (X_library[:, 5] + X_library[:, 6]) \\\n",
    "                    - X_library[:, 16] * pi_group['Eu']\n",
    "        pred_best = pred_best.reshape(-1,1)\n",
    "        r2 = r2_score(y_library, pred_best)\n",
    "        print(f\"Analytical r2_score: {round(r2, 6)}\")\n",
    "\n",
    "        pred_best = - X_library[:, 7] - X_library[:, 12]\\\n",
    "                    - X_library[:, 16] * pi_group['Eu']\n",
    "        pred_best = pred_best.reshape(-1,1)\n",
    "        r2_2 = r2_score(y_library, pred_best)\n",
    "        print(f\"Analytical r2_score (no 1/Re): {round(r2_2, 6)}\")\n",
    "        print('difference', r2-r2_2)\n",
    "        return None\n",
    "        \n",
    "    @staticmethod\n",
    "    def fit(X_library, y_library, threshold=0.002):\n",
    "        '''\n",
    "        squential threshold with dynamic threshold\n",
    "        '''\n",
    "        model = SeqReg()\n",
    "        coef, _, r2 = model.fit_dynamic_thresh(X_library, y_library, \n",
    "                        is_normalize=False, non_zero_term=3, threshold=threshold, fit_intercept=False, model_name='LR')\n",
    "        print('Fitting r2', r2)\n",
    "        return coef\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**************************************** v2-Re-100\n",
      "\u001b[31mRe: 100, coef_best: 0.01\u001b[0m\n",
      "\u001b[31mEu: 14545, coef_best: 14545.454545\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4500/4500 [00:12<00:00, 351.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analytical r2_score: 0.998459\n",
      "Analytical r2_score (no 1/Re): 0.990676\n",
      "difference 0.007783185152754779\n",
      "Fitting r2 0.9936145467269653\n",
      "\u001b[31mcoef_res: [('px', -14016.091245), ('uux', -1.043666), ('vuy', -1.043666)]\u001b[0m\n",
      "**************************************** v1-Re-170\n",
      "\u001b[31mRe: 170, coef_best: 0.005882\u001b[0m\n",
      "\u001b[31mEu: 3897, coef_best: 3897.626183\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4500/4500 [00:12<00:00, 354.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analytical r2_score: 0.999497\n",
      "Analytical r2_score (no 1/Re): 0.997588\n",
      "difference 0.0019084126014493608\n",
      "Fitting r2 0.997589885870623\n",
      "\u001b[31mcoef_res: [('px', -3896.136606), ('uux', -1.001648), ('vuy', -1.001648)]\u001b[0m\n",
      "**************************************** v3-Re-200\n",
      "\u001b[31mRe: 200, coef_best: 0.005\u001b[0m\n",
      "\u001b[31mEu: 3019, coef_best: 3019.973003\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4500/4500 [00:13<00:00, 329.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analytical r2_score: 0.99927\n",
      "Analytical r2_score (no 1/Re): 0.996591\n",
      "difference 0.002678801348929616\n",
      "Fitting r2 0.9967439478235139\n",
      "\u001b[31mcoef_res: [('px', -3004.482481), ('uux', -1.017731), ('vuy', -1.017731)]\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "config_path = 'configs/config_Euler_3cylinder_clean.yml'\n",
    "config = yaml.load(open(config_path, 'r'), Loader=yaml.FullLoader)\n",
    "case_id_list = list(config['case_id_list'].keys())\n",
    "\n",
    "def prepare_dataset(is_show=False):\n",
    "    '''\n",
    "    prepare a sets of dataset\n",
    "    '''\n",
    "    res = []\n",
    "    fit_equ = FitEqu()\n",
    "\n",
    "    for case_id in case_id_list:\n",
    "        print('*' * 40, case_id)\n",
    "        case_info = config['case_id_list'][case_id]\n",
    "        Nx, Ny = 500, 222\n",
    "        dx, dy = 0.9/Nx, 0.4/Ny\n",
    "\n",
    "        ref_params = case_info['ref_params']\n",
    "        sim_params = case_info['sim_params']\n",
    "        sim_params['niu'] = sim_params['mu'] / sim_params['rho']\n",
    "        dt = sim_params['dt']\n",
    "\n",
    "        # parameters for sparse regression\n",
    "        x_range, y_range = case_info['x_range'], case_info['y_range']\n",
    "        boundary_x, boundary_y, boundary_t = case_info['fitting']['boundary_num']  # default is 5\n",
    "        deg = case_info['fitting']['deg']               # polynomial degree default is 3\n",
    "        \n",
    "        #####################Prepare Dataset#####################\n",
    "        # generate data and save to a folder\n",
    "        U, V, P = fit_equ.load_dataset(case_id, config['dataset']['dataset_dir'])\n",
    "        U, V, P = fit_equ.select_sub_domain(U, V, P, x_range, y_range)\n",
    "        # U, V, P = fit_equ.add_noise(U, V, P, noise_level=0.01)\n",
    "        # U, V, P = fit_equ.denoise(U, V, P, case_info['SVD_mode'], is_show=True)\n",
    "        # non-dimensionlize data and dx, dy, dt\n",
    "        U, V, P, dx, dy, dt, pi_group = fit_equ.non_dimensionlize(U, V, P, dx, dy, dt, ref_params, sim_params)\n",
    "        Nx, Ny, Nt = P.shape\n",
    "        # sample points\n",
    "        points = fit_equ.sample_points(Nx, Ny, boundary_x, boundary_y, boundary_t)        \n",
    "        # plot the 1st frame \n",
    "        if is_show: fig = plt.figure(figsize=(6, 4)); plt.imshow(V[:, :, 0].T)\n",
    "        \n",
    "        #####################Prepare library#####################\n",
    "        base_library = fit_equ.cal_derivatives(\n",
    "            U, V, P, points, boundary_x, boundary_y, boundary_t, dx, dy, dt, deg)\n",
    "        X_library, y_library, names = fit_equ.parse_library(base_library, sim_params, ref_params)\n",
    "        X_library, y_library, norm_coef = fit_equ.normalization(X_library, y_library)\n",
    "        # check data\n",
    "        is_nan = fit_equ.check_library(X_library, y_library)\n",
    "        if is_nan: continue\n",
    "        df_all = pd.DataFrame(np.concatenate([X_library, y_library], axis=1))\n",
    "        # check the fitting of the target equation\n",
    "        fit_equ.check_equ(X_library, y_library, pi_group)\n",
    "        # visualization library\n",
    "        # if is_show: fig = plt.figure(); df_all.plot(); plt.xlim([0, 2000])\n",
    "\n",
    "        #####################Fit sequential sparse regression#####################\n",
    "        coef = fit_equ.fit(X_library, y_library)\n",
    "        # coef = np.multiply(norm_coef, coef)\n",
    "        # print('rescale coef', coef)\n",
    "        coef_res = [(each[0], round(each[1], 6)) for each in list(zip(names, coef.tolist())) if abs(each[1]) > 1e-3]\n",
    "        coef_res = sorted(coef_res, key=lambda x: abs(x[1]), reverse=True)\n",
    "        print(colored(f'coef_res: {coef_res}', 'red'))\n",
    "        # add to the total results\n",
    "        res.append([sim_params['niu'], ref_params['l_ref'], ref_params['v_ref'], ref_params['p_ref'],\n",
    "                sim_params['rho'], sim_params['g'], pi_group['Re'], pi_group['Eu'], pi_group['Fr']] + coef.tolist())\n",
    "        \n",
    "    df = pd.DataFrame(res)\n",
    "    return df\n",
    "\n",
    "\n",
    "df = prepare_dataset()\n",
    "df.to_csv('../dataset/dataset_NS_flow3D.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('../dataset/dataset_NS_flow3D.csv')\n",
    "df = df.rename(columns={\n",
    "    '0': 'nu',\n",
    "    '1': 'diameter',\n",
    "    '2': 'v_init',\n",
    "    '3': 'p_init',\n",
    "    '4': 'density',\n",
    "    '5': 'g',\n",
    "    '6': 'Re',\n",
    "    '7': 'Eu',\n",
    "    '8': 'Fr',\n",
    "    })\n",
    "df['mu'] = df['nu'] * df['density']\n",
    "df['1/Re'] = 1/ df['Re']\n",
    "# print(df.head())\n",
    "df.to_csv('../dataset/dataset_NS_flow3D.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Eu</th>\n",
       "      <th>25</th>\n",
       "      <th>21</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>14545.454545</td>\n",
       "      <td>-14016.091245</td>\n",
       "      <td>-1.043666</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3897.626183</td>\n",
       "      <td>-3896.136606</td>\n",
       "      <td>-1.001648</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3019.973003</td>\n",
       "      <td>-3004.482481</td>\n",
       "      <td>-1.017731</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             Eu            25        21\n",
       "0  14545.454545 -14016.091245 -1.043666\n",
       "1   3897.626183  -3896.136606 -1.001648\n",
       "2   3019.973003  -3004.482481 -1.017731"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[['Eu', '25', '21']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-10.203738348681851"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r2_score(df['1/Re'], df['14'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DimensionlessLearning(object):\n",
    "    '''\n",
    "    Indentify the explicit form one coefficient using dimensionless learning\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, df, input_list, output_coef, dimension_info, basis_list):\n",
    "        super(DimensionlessLearning, self).__init__()\n",
    "        self.df = df\n",
    "        self.input_list = input_list\n",
    "        self.output_coef = output_coef\n",
    "        self.X, self.y = self.prepare_dataset()\n",
    "        self.dimension_info, self.basis_list = dimension_info, basis_list\n",
    "        self.basis1_in, self.basis2_in = self.prepare_dimension()\n",
    "\n",
    "    def prepare_dataset(self):\n",
    "        '''\n",
    "        prepare the input and output data\n",
    "        '''\n",
    "        X = self.df[self.input_list].to_numpy()\n",
    "        y = self.df[self.output_coef].to_numpy().reshape(-1, 1)\n",
    "        return X, y\n",
    "        \n",
    "    def prepare_dimension(self):\n",
    "        '''\n",
    "        parse dimension for input and output\n",
    "        '''\n",
    "        basis1_in, basis2_in = self.basis_list[0], self.basis_list[1]\n",
    "        return basis1_in, basis2_in\n",
    "\n",
    "    def fetch_coef_pi(self, coef):\n",
    "        '''\n",
    "        parse the combined weights for the input\n",
    "        '''\n",
    "        coef_pi = coef[0] * self.basis1_in + coef[1] * self.basis2_in\n",
    "        return coef_pi\n",
    "        \n",
    "    def check_dimension(self, coef):\n",
    "        '''\n",
    "        check whether the basis vectors can formulated as the D_out\n",
    "        '''\n",
    "        coef_pi = self.fetch_coef_pi(coef)\n",
    "        # print('[check] coef_pi: \\n', coef_pi)\n",
    "        target_D_out = np.dot(self.dimension_info[0], coef_pi)\n",
    "        # print('[check] target_D_out: \\n', target_D_out)\n",
    "        assert np.array_equal(target_D_out, self.dimension_info[1]), 'Wrong target_D_out!'\n",
    "\n",
    "    def fit_pattern_search(self, seed):\n",
    "        '''\n",
    "        pattern search\n",
    "        '''\n",
    "        def get_coordinates(a, b, delta):\n",
    "            '''\n",
    "            Build a list to store all possible coordiantes\n",
    "            '''\n",
    "            coord_all = []\n",
    "            for a_ in [a-delta, a, a+delta]:\n",
    "                for b_ in [b-delta, b, b+delta]:\n",
    "                    if [a_, b_] != [a, b]:\n",
    "                        coord_all.append([a_, b_])\n",
    "            return coord_all\n",
    "        \n",
    "        def opt(coef):\n",
    "            '''\n",
    "            fit a linear regression\n",
    "            '''\n",
    "            coef_pi = self.fetch_coef_pi(coef)\n",
    "            pi_in = np.prod(np.power(self.X, coef_pi.reshape(-1,)), axis=1).reshape(-1, 1)\n",
    "            reg =LinearRegression(fit_intercept=False)\n",
    "            reg.fit(pi_in, self.y)\n",
    "            y_pred = reg.predict(pi_in)\n",
    "            r2 = r2_score(self.y, y_pred)\n",
    "            return r2, coef_pi, reg.coef_\n",
    "\n",
    "        np.random.seed(seed)\n",
    "        res, break_points = [], []\n",
    "        a = np.random.choice(np.linspace(-2, 2, 9), 1)[0]  # [-2, 2] delta=0.5\n",
    "        b = np.random.choice(np.linspace(-2, 2, 9), 1)[0]  # [-2, 2] delta=0.5\n",
    "        # a, b = -1, 0\n",
    "        coef = np.array([a, b]).reshape(-1, 1)\n",
    "\n",
    "        iter_num, max_iter, delta = 0, 10, 0.5\n",
    "        while iter_num < max_iter:\n",
    "            candidate_coord = get_coordinates(a, b, delta)\n",
    "            r2_center, reg_coef_center, coef_w_center = opt(coef)\n",
    "            # print('r2_center', round(r2_center, 2), 'reg_coef_center', [round(each, 2) for each in list(reg_coef_center.reshape(-1,))])\n",
    "            # print('coef_w_center', coef_w_center)\n",
    "\n",
    "            if r2_center < 0.2:\n",
    "                break_points.append([a, b])\n",
    "                break\n",
    "            \n",
    "            r2_bounds_val = []\n",
    "            for [a_, b_] in candidate_coord:\n",
    "                coef_temp = np.array([a_, b_]).reshape(-1, 1)\n",
    "                r2_bound, reg_coef_bound, coef_w_bound = opt(coef_temp)\n",
    "                r2_bounds_val.append(r2_bound)\n",
    "\n",
    "            # sort r2 from high to low\n",
    "            highest_index = np.argsort(r2_bounds_val)[::-1][0]\n",
    "            iter_num += 1\n",
    "\n",
    "            # udpate the center coordiantes when the R2 in the neighborhood is higher\n",
    "            if r2_center < r2_bounds_val[highest_index]:\n",
    "                [a, b] = candidate_coord[highest_index]\n",
    "                coef = np.array([a, b]).reshape(-1, 1)\n",
    "                coef_pi = self.fetch_coef_pi(coef)\n",
    "                res_info = {'a': a, 'b': b, 'r2_center': round(r2_bounds_val[highest_index], 4)}\n",
    "                # print('update', res_info)\n",
    "                res.append(res_info)\n",
    "            else:\n",
    "                break\n",
    "        \n",
    "        coef_pi = self.fetch_coef_pi(coef)\n",
    "        r2, reg_coef_final, coef_w_final = opt(coef)\n",
    "        return r2, reg_coef_final, coef_w_final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "final r2 0.9976327599354018 [-2.5  2.5 -1.5  0.5  2. ] [[-0.]]\n",
      "final r2 0.9976327599354018 [-2.5  2.5 -1.5  0.5  2. ] [[-0.]]\n",
      "final r2 0.9996567905264077 [ 0.  0. -2. -1.  1.] [[-0.96713434]]\n",
      "final r2 0.9996567905264077 [ 0.  0. -2. -1.  1.] [[-0.96713434]]\n",
      "final r2 0.9996567905264077 [ 0.  0. -2. -1.  1.] [[-0.96713434]]\n",
      "final r2 0.9996567905264077 [ 0.  0. -2. -1.  1.] [[-0.96713434]]\n",
      "final r2 0.9976327599354018 [-2.5  2.5 -1.5  0.5  2. ] [[-0.]]\n"
     ]
    }
   ],
   "source": [
    "def recover_coef1(seed):\n",
    "    input_list = ['mu', 'diameter', 'v_init', 'density', 'p_init']\n",
    "    output_coef = '25'  # for Eu\n",
    "\n",
    "    D_in = np.mat('-1, -1, 1; 1, 0, 0; 1, -1, 0; -3, 0, 1; -1, -2, 1').T\n",
    "    D_out = np.mat('0;, 0; 0')\n",
    "    dimension_info = [D_in, D_out]\n",
    "\n",
    "    basis1_in = np.array([-1, 1, 1, 1, 0]).reshape(-1, 1)\n",
    "    basis2_in = np.array([-1, 1, -1, 0, 1]).reshape(-1, 1)\n",
    "    basis_list = [basis1_in, basis2_in]\n",
    "    \n",
    "    dimensionless_learning = DimensionlessLearning(\n",
    "        df, input_list, output_coef, dimension_info, basis_list)\n",
    "    # dimensionless_learning.check_dimension(coef=[-1, 1])\n",
    "\n",
    "    # pattern search\n",
    "    r2, coef, coef_w = dimensionless_learning.fit_pattern_search(seed=seed)\n",
    "    if r2 > 0.8:\n",
    "        print('final r2', r2, coef.flatten(), coef_w)\n",
    "\n",
    "\n",
    "for i in range(10):\n",
    "    recover_coef1(seed=i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "xie",
   "language": "python",
   "name": "python3"
  },
  "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.8.8 (default, Feb 24 2021, 21:46:12) \n[GCC 7.3.0]"
  },
  "vscode": {
   "interpreter": {
    "hash": "d4e05679d6a31593988b5489bd865ca5a262d5481eb276cba948c252ea078b95"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}