{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solve viscous Burgers equation with SWIM-ODE:\n",
    "In this script we show how SWIM-ODE uses resampling of collocation points and data-dependant sampling of weights and biases to solve Burgers equation accurately"
   ]
  },
  {
   "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",
    "import scipy\n",
    "from sklearn.metrics import mean_squared_error\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 SWIM-ODE network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\PhD\\Work\\Programming\\swim_pde_paper\\swimpde-paper\\swimpde\\examples\\burgers\\../../src\\swimpde\\Solver\\burgers_solver.py:144: RuntimeWarning: overflow encountered in multiply\n",
      "  non_lin_term =  t1 * t2\n",
      "c:\\PhD\\Work\\Programming\\swim_pde_paper\\swimpde-paper\\swimpde\\examples\\burgers\\../../src\\swimpde\\Solver\\burgers_solver.py:153: RuntimeWarning: invalid value encountered in matmul\n",
      "  c_t = self._A_boundary_inv @ np.concatenate([rhs, self.scale_boundary_correction * boundary_correction])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rmse_swim, re_swim\n",
      "0.0006080909977263369 0.0009898321100579146\n",
      "rmse_swim, re_swim\n",
      "4.064199834080494e-06 6.61558140558983e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\PhD\\Work\\Programming\\swim_pde_paper\\swimpde-paper\\swimpde\\examples\\burgers\\../../src\\swimpde\\Solver\\burgers_solver.py:144: RuntimeWarning: overflow encountered in multiply\n",
      "  non_lin_term =  t1 * t2\n",
      "c:\\PhD\\Work\\Programming\\swim_pde_paper\\swimpde-paper\\swimpde\\examples\\burgers\\../../src\\swimpde\\Solver\\burgers_solver.py:153: RuntimeWarning: invalid value encountered in matmul\n",
      "  c_t = self._A_boundary_inv @ np.concatenate([rhs, self.scale_boundary_correction * boundary_correction])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rmse_swim, re_swim\n",
      "3.0048800977784118e-06 4.8912528203444845e-06\n"
     ]
    }
   ],
   "source": [
    "# Hyper-parameters\n",
    "n_sample = 3000 # No. of sampling points (for computing gradients)\n",
    "n_col = 800 # No. of collocation points (to be re-sampled)\n",
    "width = 450 # Width\n",
    "reg_const = 1e-7 # Regularization constant\n",
    "svd_cutoff = 1e-9 # SVD threshold\n",
    "seeds = [1,2,3] # Seeds (to compute mean errors)\n",
    "time_blocks = 3 # Number of time-blocks for smapling\n",
    "info = [] # List to store errors and time measurements\n",
    "\n",
    "# Define prob. distribution for (re)-sampling collocation points\n",
    "def collocation_points_probabilities(df_dx):\n",
    "    gradients = np.abs(df_dx)\n",
    "    gradients = gradients + 0.01 * 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",
    "# Store errors and time required for each seed\n",
    "rel_err_swim = np.ones((len(seeds), ))\n",
    "time_swim = np.ones((len(seeds), ))\n",
    "rmse_swim = np.ones((len(seeds), ))\n",
    "j = 0\n",
    "for seed in seeds: # Run over 3 seeds                                    \n",
    "    # SWIM network ansatz\n",
    "    ansatz_swim = BasicAnsatz(\n",
    "        n_neurons=width,\n",
    "        activation=\"tanh\",\n",
    "        random_state=seed,\n",
    "        regularization_scale=reg_const,\n",
    "        parameter_sampler = \"tanh\" \n",
    "    )\n",
    "    # Burgers equation solver\n",
    "    burgers_solver_swim = BurgersSolver(\n",
    "        domain=domain, \n",
    "        ansatz=ansatz_swim,\n",
    "        u0=u0,\n",
    "        boundary_condition=boundary_condition,\n",
    "        forcing=forcing,\n",
    "        regularization_scale=reg_const,\n",
    "        c=(0.01/np.pi)\n",
    "    )\n",
    "    # SWIM fit\n",
    "    t_swim_start = time.time()\n",
    "    sol_swim, solver_status_swim = burgers_solver_swim.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_swim_stop = time.time()\n",
    "    time_swim[j] = t_swim_stop - t_swim_start\n",
    "    # Evaluate SWIM\n",
    "    u_swim = (burgers_solver_swim.evaluate_blocks(x_eval= x_eval, t_eval = t_eval, time_blocks = time_blocks, solver_status = solver_status_swim)).T\n",
    "    # Compute metrics\n",
    "    mse_swim = mean_squared_error(u_true, u_swim.flatten()[:,None])  # mean squared error\n",
    "    rmse_swim[j] = np.sqrt(mse_swim)  # Root Mean Squared Error                  \n",
    "    rel_err_swim[j] = np.linalg.norm(u_true-u_swim.flatten()[:,None],2)/np.linalg.norm(u_true,2) # Relative l-2 error\n",
    "    print(\"rmse_swim, re_swim\")\n",
    "    print(rmse_swim[j], rel_err_swim[j])\n",
    "    j += 1\n",
    "\n",
    "# RMSE and relative l-2 errors (mean and std) and time measurements\n",
    "info.append(np.mean(time_swim))\n",
    "info.append(np.mean(rmse_swim))\n",
    "info.append(np.std(rmse_swim))\n",
    "info.append(np.mean(rel_err_swim))\n",
    "info.append(np.std(rel_err_swim))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Errors and time measurements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Burgers equation: Errors and time measurements using ODE-SWIM\n",
      "training time for swim =  360.7527981599172\n",
      "rmse swim =  0.00020505335921939863 +- 0.00028499097538871066\n",
      "rel error swim =  0.0003337796480946163 +- 0.0004638996787836397\n"
     ]
    }
   ],
   "source": [
    "# Print errors and time measurements for Burgers with ODE-SWIM\n",
    "res = np.vstack(info).reshape(-1)\n",
    "\n",
    "# Burgers time measurements\n",
    "print('Burgers equation: Errors and time measurements using ODE-SWIM')\n",
    "print('training time for swim = ', res[-5])\n",
    "print('rmse swim = ', res[-4], '+-', res[-3])\n",
    "print('rel error swim = ', res[-2], '+-', res[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error plots and comparison of solutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative L2 error on the entire spatio-temporal domain: 4.891253e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\chinm\\AppData\\Local\\Temp\\ipykernel_32156\\1003889938.py:82: UserWarning: The figure layout has changed to tight\n",
      "  fig.tight_layout()\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32708 characters]",
      "text/plain": [
       "<Figure size 600x300 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 34460 characters]",
      "text/plain": [
       "<Figure size 600x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# True and model solutions\n",
    "u_true_test = u_exact\n",
    "u_model_test = u_swim\n",
    "\n",
    "# Relative l-2 error\n",
    "error_u = np.linalg.norm(u_true-u_model_test.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",
    "fontsize = 14\n",
    "fig, ax = plt.subplots(1, 3, figsize=(6, 3), constrained_layout=True)\n",
    "extent = [0 , 1, np.min(x_space), np.max(x_space)]\n",
    "aspect = 0.3\n",
    "sol_img1 = ax[0].imshow(u_true_test.T, extent=extent, origin='lower', aspect=aspect)#, vmin=vlim[0], vmax=vlim[1]\n",
    "sol_img2 = ax[1].imshow(u_model_test.T, extent=extent, origin='lower', aspect=aspect)#, vmin=vlim[0], vmax=vlim[1]\n",
    "error_img = ax[2].imshow(abs(u_model_test - u_true_test).T,  extent=extent, origin='lower', aspect=aspect)\n",
    "\n",
    "ax[0].set_xlabel('t', fontsize = fontsize)\n",
    "ax[1].set_xlabel('t', fontsize = fontsize)\n",
    "ax[2].set_xlabel('t', fontsize = fontsize)\n",
    "ax[0].set_ylabel('x', fontsize = fontsize)\n",
    "\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",
    "\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",
    "\n",
    "cbar_true = fig.colorbar(sol_img1, ax=ax[0], location='bottom')\n",
    "cbar_swim = 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_swim = ticker.MaxNLocator(nbins=2)\n",
    "cbar_swim.locator = tick_locator_swim\n",
    "cbar_swim.update_ticks()\n",
    "\n",
    "tick_locator_true = ticker.MaxNLocator(nbins=2)\n",
    "cbar_true.locator = tick_locator_true\n",
    "cbar_true.update_ticks()\n",
    "ax[0].set_title('Ground truth')\n",
    "ax[1].set_title('SWIM-ODE')\n",
    "ax[2].set_title('Absolute error')\n",
    "fig.savefig(\"burgers_1.pdf\")\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(6,3), constrained_layout=True)\n",
    "plt.tick_params(axis='both', labelsize=fontsize)\n",
    "ax[0].plot(x_eval,u_exact[25,:], 'b-', linewidth = 2, label = 'Ground truth')       \n",
    "ax[0].plot(x_eval,u_model_test[25,:], 'r--', linewidth = 2, label = 'SWIM-ODE (resampling)')\n",
    "ax[0].set_xlabel('$x$', fontsize = fontsize)\n",
    "ax[0].set_ylabel('$u(t,x)$', fontsize = fontsize)    \n",
    "ax[0].set_title('$t = 0.25$', fontsize = fontsize)\n",
    "ax[0].axis('square')\n",
    "ax[0].set_xlim([-1.1,1.1])\n",
    "ax[0].set_ylim([-1.1,1.1])\n",
    "plt.tick_params(axis='both', labelsize=fontsize)\n",
    " \n",
    "ax[1].plot(x_eval,u_exact[50,:], 'b-', linewidth = 2, label = 'Ground truth')       \n",
    "ax[1].plot(x_eval,u_model_test[50,:], 'r--', linewidth = 2, label = 'SWIM-ODE (resampling)')\n",
    "ax[1].set_xlabel('$x$', fontsize = fontsize)\n",
    "ax[1].set_ylabel('$u(t,x)$', fontsize = fontsize)\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 = fontsize)\n",
    "    \n",
    "ax[2].plot(x_eval,u_exact[75,:], 'b-', linewidth = 2, label = 'Ground truth')       \n",
    "ax[2].plot(x_eval,u_model_test[75,:], 'r--', linewidth = 2, label = 'SWIM-ODE (resampling)')\n",
    "ax[2].set_xlabel('$x$', fontsize = 14)\n",
    "ax[2].set_ylabel('$u(t,x)$', fontsize = fontsize)\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 = fontsize)\n",
    "\n",
    "# Create a single legend for all plots\n",
    "fig.legend(*ax[1].get_legend_handles_labels(),loc='upper center', ncol=2, fontsize = fontsize, frameon=False)\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"burgers_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
}
