{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solve viscous Burgers equation with ODE-ELM to show that it fails!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../../')\n",
    "sys.path.append('../../src')\n",
    "from swimpde import Domain\n",
    "from swimpde import BasicAnsatz\n",
    "from swimpde import BurgersSolver\n",
    "import numpy as np\n",
    "from sklearn.metrics import mean_squared_error\n",
    "import scipy.io\n",
    "from matplotlib import ticker\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['image.cmap'] = 'jet'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data and reference solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and visualize data\n",
    "data = scipy.io.loadmat('../../data/burgers_shock.mat')\n",
    "t_eval = data['t'].flatten()[:,None]\n",
    "x_eval = data['x'].flatten()[:,None]\n",
    "u_exact = np.real(data['usol']).T\n",
    "X, T = np.meshgrid(x_eval,t_eval)\n",
    "X_ = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "\n",
    "# Set ground truth\n",
    "u_true = u_exact.flatten()[:,None]              "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial condition\n",
    "def u0(x):\n",
    "    return -1 * np.sin(np.pi * x)\n",
    "\n",
    "# forcing\n",
    "def forcing(x, t):\n",
    "    return np.zeros(x.shape[0])\n",
    "\n",
    "# boundary condition\n",
    "boundary_condition = \"zero dirichlet\" # \n",
    "\n",
    "# Domain information and spacial points for the first time-block\n",
    "n_points_1d = 4000 # No. of points in space\n",
    "x_lim = [-1, 1] # Domain range\n",
    "\n",
    "# Interior points\n",
    "rng = np.random.default_rng(seed=123)\n",
    "x_space = rng.uniform(x_lim[0], x_lim[1], n_points_1d).reshape((-1, 1)) \n",
    "x_space_inner = x_space[1:-1]\n",
    "interior_points = x_space_inner\n",
    "\n",
    "# Boundary points (excluding corners)\n",
    "left = x_lim[0]\n",
    "right = x_lim[1]\n",
    "boundary_points = np.row_stack([left, right])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train and evaluate ODE-ELM network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rmse_elm, re_elm\n",
      "0.15175027104385325 0.15175027104385325 0.24701449544694037 0.24701449544694037\n",
      "rmse_elm, re_elm\n",
      "0.15173057852918842 0.15173057852918842 0.24698244056795698 0.24698244056795698\n",
      "rmse_elm, re_elm\n",
      "0.15243588299644448 0.15243588299644448 0.24813051381960463 0.24813051381960463\n",
      "[6.187496344248454, 0.15197224418982871, 0.00032794070199099765, 0.24737581661150065, 0.0005338119430140186]\n"
     ]
    }
   ],
   "source": [
    "# Hyper-parameters\n",
    "n_sample = 6000 # No. of sampling points (for computing gradients)\n",
    "n_col = 3000 # No. of collocation points (to be re-sampled)\n",
    "width = 2000 # Width\n",
    "reg_const = 1e-7 # Regularization constant\n",
    "svd_cutoff = 1e-10 # SVD threshold\n",
    "seeds = [1, 2, 3] # Seeds (to compute mean errors)\n",
    "time_blocks = 1 # Number of time-blocks for smapling\n",
    "info = [] # List to store errors and time measurements\n",
    "\n",
    "# Compute prob. distribution for (re)-sampling collocation points\n",
    "def collocation_points_probabilities(df_dx):\n",
    "    gradients = np.abs(df_dx)\n",
    "    gradients = gradients + 0.05 * np.max(gradients)\n",
    "    return gradients/np.sum(gradients)\n",
    "\n",
    "# Points where gradient of the solution at the end of a time-block is computed\n",
    "sample_test_points = np.sort(rng.uniform(x_lim[0] + 1e-4, x_lim[1] - 1e-4, n_sample)).reshape((-1, 1)) # This does not include boundary points\n",
    "\n",
    "# Domain\n",
    "domain = Domain(\n",
    "    interior_points=interior_points,\n",
    "    boundary_points=boundary_points,\n",
    "    sample_points = sample_test_points\n",
    ")\n",
    "\n",
    "# Parameter sampler for sampling of weights and biases using data-agnostic distribution\n",
    "def parameter_sampler_uniform(x, y, rng):\n",
    "    \"\"\"\n",
    "        returns: weights, biases, idx_from, idx_to\n",
    "    \"\"\"\n",
    "    n_dim = 1\n",
    "    x_left = -3\n",
    "    x_right = 3\n",
    "    # Sample weights from a normal distribution\n",
    "    weights = rng.normal(size=(n_dim, width)) #rng.uniform(x_left, x_right, size=(n_dim, n_OBF))\n",
    "    # Sample biases from a uniform distribution\n",
    "    biases = rng.uniform(x_left, x_right, size=(1, width))\n",
    "    idx_from = np.arange(width)\n",
    "    idx_to = np.arange(width)\n",
    "    return weights, biases, idx_from, idx_to\n",
    "\n",
    "# Loop over different seeds\n",
    "rel_err_elm = np.ones((len(seeds), ))\n",
    "time_elm = np.ones((len(seeds), ))\n",
    "rmse_elm = np.ones((len(seeds), ))\n",
    "j = 0\n",
    "\n",
    "# Loop over different seeds\n",
    "for seed in seeds: # Run for 3 seeds \n",
    "    # ODE-ELM network ansatz                     \n",
    "    ansatz_elm = BasicAnsatz(\n",
    "        n_neurons=width,\n",
    "        activation=\"tanh\",\n",
    "        random_state=seed,\n",
    "        regularization_scale=reg_const,\n",
    "        parameter_sampler = parameter_sampler_uniform  \n",
    "    )\n",
    "    # Burger solver\n",
    "    burgers_solver_elm = BurgersSolver(\n",
    "        domain=domain, \n",
    "        ansatz=ansatz_elm,\n",
    "        u0=u0,\n",
    "        boundary_condition=boundary_condition,\n",
    "        forcing=forcing,\n",
    "        regularization_scale=reg_const,\n",
    "        c=(0.01/np.pi)\n",
    "    )\n",
    "    # ELM fit\n",
    "    t_elm_start = time.time()\n",
    "    sol_elm, solver_status_elm = burgers_solver_elm.fit_time_blocks(t_span=[0, np.max(t_eval)], rtol=1e-8, atol=1e-8, svd_cutoff=svd_cutoff, time_blocks=time_blocks, prob_distr_resampling = collocation_points_probabilities, n_col=n_col, outer_basis=False);\n",
    "    t_elm_stop = time.time()\n",
    "    time_elm[j] = t_elm_stop - t_elm_start\n",
    "    # Evaluate ELM-ODE predictions\n",
    "    u_elm = (burgers_solver_elm.evaluate_blocks(x_eval= x_eval, t_eval = t_eval, time_blocks = time_blocks, solver_status = solver_status_elm)).T\n",
    "    # Compute metrics\n",
    "    mse_elm = mean_squared_error(u_true, u_elm.flatten()[:,None])  # mean squared error\n",
    "    rmse_elm[j] = np.sqrt(mse_elm)  # Root Mean Squared Error\n",
    "    rel_err_elm[j] = np.linalg.norm(u_true-u_elm.flatten()[:,None], 2)/np.linalg.norm(u_true,2)\n",
    "    print(\"rmse_elm, re_elm\")\n",
    "    print(rmse_elm[j], rmse_elm[j], rel_err_elm[j], rel_err_elm[j])\n",
    "    j += 1\n",
    "\n",
    "info.append(np.mean(time_elm))\n",
    "info.append(np.mean(rmse_elm))\n",
    "info.append(np.std(rmse_elm))\n",
    "info.append(np.mean(rel_err_elm))\n",
    "info.append(np.std(rel_err_elm))\n",
    "print(info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Burgers equation: Errors and time measurements using ode-elm\n",
      "training time for ode-elm =  6.187496344248454\n",
      "rmse ode-elm =  0.15197224418982871 +- 0.00032794070199099765\n",
      "rel error ode-elm =  0.24737581661150065 +- 0.0005338119430140186\n"
     ]
    }
   ],
   "source": [
    "# Print errors and time measurements for Burgers with ODE-ELM\n",
    "res = np.vstack(info).reshape(-1)\n",
    "\n",
    "# Burgers time measurements\n",
    "print('Burgers equation: Errors and time measurements using ode-elm')\n",
    "print('training time for ode-elm = ', res[-5])\n",
    "print('rmse ode-elm = ', res[-4], '+-', res[-3])\n",
    "print('rel error ode-elm = ', res[-2], '+-', res[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative L2 error on the entire spatio-temporal domain: 2.481305e-01\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\chinm\\AppData\\Local\\Temp\\ipykernel_14040\\841674867.py:69: UserWarning: The figure layout has changed to tight\n",
      "  fig.tight_layout()\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34788 characters]",
      "text/plain": [
       "<Figure size 500x400 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 39400 characters]",
      "text/plain": [
       "<Figure size 600x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# True and model solutions\n",
    "error_u = np.linalg.norm(u_true-u_elm.flatten()[:,None]  ,2)/np.linalg.norm(u_true,2)\n",
    "print('Relative L2 error on the entire spatio-temporal domain: %e' % (error_u))\n",
    "\n",
    "# visualize the solution\n",
    "fig, ax = plt.subplots(1, 3, figsize=(5, 4), constrained_layout=True)\n",
    "fontsize = 14\n",
    "extent = [0 , 1, np.min(x_space), np.max(x_space)]\n",
    "aspect = 0.3\n",
    "sol_img1 = ax[0].imshow(u_exact.T, extent=extent, origin='lower', aspect=aspect)#, vmin=vlim[0], vmax=vlim[1]\n",
    "sol_img2 = ax[1].imshow(u_elm.T, extent=extent, origin='lower', aspect=aspect)#, vmin=vlim[0], vmax=vlim[1]\n",
    "error_img = ax[2].imshow(abs(u_elm - u_exact).T,  extent=extent, origin='lower', aspect=aspect)\n",
    "ax[0].axvline(x=0.25, color='k', linestyle='--', linewidth=2)\n",
    "ax[0].axvline(x=0.5, color='k', linestyle='--', linewidth=2)\n",
    "ax[0].axvline(x=0.75, color='k', linestyle='--', linewidth=2)\n",
    "ax[1].axvline(x=0.33, color='gray', linestyle='dotted', linewidth=3)\n",
    "ax[1].axvline(x=0.66, color='gray', linestyle='dotted', linewidth=3)\n",
    "cbar_true = fig.colorbar(sol_img1, ax=ax[0], location='bottom')\n",
    "cbar_elm = fig.colorbar(sol_img2, ax=ax[1], location='bottom')\n",
    "cbar_err = fig.colorbar(error_img, ax=ax[2], location='bottom', format='%.0e', fraction=0.046)\n",
    "\n",
    "tick_locator_err = ticker.MaxNLocator(nbins=2)\n",
    "cbar_err.locator = tick_locator_err\n",
    "cbar_err.update_ticks()\n",
    "\n",
    "tick_locator_elm = ticker.MaxNLocator(nbins=2)\n",
    "cbar_elm.locator = tick_locator_elm\n",
    "cbar_elm.update_ticks()\n",
    "\n",
    "tick_locator_true = ticker.MaxNLocator(nbins=2)\n",
    "cbar_true.locator = tick_locator_true\n",
    "cbar_true.update_ticks()\n",
    "\n",
    "ax[0].set_title('Ground truth')\n",
    "ax[1].set_title('ELM-ODE')\n",
    "ax[2].set_title('Absolute error')\n",
    "fig.savefig(\"burgers_elm_ode_1.pdf\")\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(6,3), constrained_layout=True)\n",
    "ax[0].plot(x_eval,u_exact[25,:], 'b-', linewidth = 2, label = 'Ground truth')       \n",
    "ax[0].plot(x_eval,u_elm[25,:], 'r--', linewidth = 2, label = 'ELM-ODE (resampling)')\n",
    "ax[0].set_xlabel('$x$', fontsize = 14)\n",
    "ax[0].set_ylabel('$u(t,x)$', fontsize = 14)    \n",
    "ax[0].set_title('$t = 0.25$', fontsize = 14)\n",
    "ax[0].axis('square')\n",
    "ax[0].set_xlim([-1.1,1.1])\n",
    "ax[0].set_ylim([-1.1,1.1])\n",
    "    \n",
    "ax[1].plot(x_eval,u_exact[50,:], 'b-', linewidth = 2, label = 'Ground truth')       \n",
    "ax[1].plot(x_eval,u_elm[50,:], 'r--', linewidth = 2, label = 'ELM-ODE (resampling)')\n",
    "ax[1].set_xlabel('$x$', fontsize = 14)\n",
    "ax[1].set_ylabel('$u(t,x)$', fontsize = 14)\n",
    "ax[1].axis('square')\n",
    "ax[1].set_xlim([-1.1,1.1])\n",
    "ax[1].set_ylim([-1.1,1.1])\n",
    "ax[1].set_title('$t = 0.50$', fontsize = 14)\n",
    "    \n",
    "ax[2].plot(x_eval,u_exact[75,:], 'b-', linewidth = 2, label = 'Ground truth')       \n",
    "ax[2].plot(x_eval,u_elm[75,:], 'r--', linewidth = 2, label = 'ELM-ODE (resampling)')\n",
    "ax[2].set_xlabel('$x$', fontsize = 14)\n",
    "ax[2].set_ylabel('$u(t,x)$', fontsize = 14)\n",
    "ax[2].axis('square')\n",
    "ax[2].set_xlim([-1.1,1.1])\n",
    "ax[2].set_ylim([-1.1,1.1])    \n",
    "ax[2].set_title('$t = 0.75$', fontsize = 14)\n",
    "\n",
    "# Create a single legend for all plots\n",
    "fig.legend(*ax[1].get_legend_handles_labels(),loc='upper center', ncol=2, fontsize = 12, frameon=False)\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"burgers_elm_ode_2.pdf\", bbox_inches='tight')\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
