{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving advection equation using ELM-ODE and SWIM-ODE for different values of flow velocities\n",
    "In this script, we see how the SWIM-ODE and ELM-ODE performs for different values of flow velocities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import sys\n",
    "sys.path.append('../../')\n",
    "sys.path.append('../../src')\n",
    "from swimpde import Domain\n",
    "from swimpde import BoundaryCompliantAnsatz\n",
    "from swimpde import AdvectionSolver\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import mean_squared_error\n",
    "np.random.seed(2)\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create equally spaced points in a square\n",
    "n_points_1d = 400 \n",
    "x_lim = [0, 2 * np.pi]\n",
    "\n",
    "# coordinates of boundary points (excluding corners)\n",
    "left = x_lim[0]\n",
    "right = x_lim[1]\n",
    "boundary_points = np.row_stack([left, right])\n",
    "\n",
    "# initial condition\n",
    "def u0(x):\n",
    "    return np.sin(x)\n",
    "\n",
    "# forcing\n",
    "def forcing(x, t):\n",
    "    return np.zeros(x.shape[0])\n",
    "\n",
    "# boundary condition\n",
    "boundary_condition = \"periodic\" #strict\n",
    "\n",
    "# Analytical solution\n",
    "def analytical_sol(x, t, beta):\n",
    "    return np.sin(x - beta * t)\n",
    "\n",
    "# Test data\n",
    "t_eval = np.linspace(0, 1, 100) # time domain\n",
    "x_space_test = np.linspace(x_lim[0], x_lim[-1], 256).reshape((-1, 1)) # space domain\n",
    "xx_test, yy_test = np.meshgrid(x_space_test, t_eval)\n",
    "X_test = np.hstack((xx_test.reshape(-1, 1), yy_test.reshape(-1, 1))) # Uniform test points\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit and evaluate SWIM-ODE and ELM-ODE for good hyper-parameters "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.033757527669270836, 0.043172359466552734, 6.217344519192376e-10, 2.4616337261142887e-10, 4.376065445360843e-10, 1.3672746304091418e-10, 8.809876555990293e-10, 3.488095148368321e-10, 6.200813909468158e-10, 1.9374060219533226e-10]\n",
      "[0.03362154960632324, 0.03935082753499349, 6.246828525178684e-10, 2.446835408742689e-10, 4.37025971573426e-10, 1.363557135024605e-10, 8.85165376400437e-10, 3.467125721219772e-10, 6.192586479128512e-10, 1.9321381398619897e-10]\n",
      "[0.03520782788594564, 0.040271759033203125, 6.75304471080849e-10, 2.312712505288523e-10, 5.136094542118958e-10, 1.84073482031596e-10, 9.568829609843076e-10, 3.2770332268405387e-10, 7.277667428281751e-10, 2.6082572538453967e-10]\n",
      "[0.033083438873291016, 0.04036959012349447, 1.7918848947462148e-09, 2.4852940211331203e-10, 2.6204806124823794e-09, 1.1500385541036766e-09, 2.5363530226109957e-09, 3.517850404933965e-10, 3.7092024948982767e-09, 1.6278410356448487e-09]\n",
      "[0.040961901346842446, 0.043559869130452476, 4.16722760872261e-09, 1.2611125708197152e-09, 1.0116718670261828e-08, 5.811514371521282e-09, 5.893949327483123e-09, 1.7836639335718467e-09, 1.430865621021439e-08, 8.219558526150662e-09]\n",
      "[0.0680088996887207, 0.06788341204325359, 1.3570838794486394e-08, 6.0210295403999214e-09, 3.7263429547958334e-08, 2.2532724633715202e-08, 1.91917953444009e-08, 8.514902317526264e-09, 5.2697709002738454e-08, 3.1865638246155986e-08]\n",
      "[0.12776931126912436, 0.12627657254536948, 3.309170044130356e-08, 1.5574359335397154e-08, 9.217531990270805e-08, 5.6157623520084134e-08, 4.6799160634120045e-08, 2.2025672135029356e-08, 1.3035678267061905e-07, 7.941960095427864e-08]\n",
      "[0.4315965970357259, 0.43122514088948566, 1.6405340060529916e-07, 7.889599538799764e-08, 4.5885777790778104e-07, 2.8056912624094756e-07, 2.320094252675112e-07, 1.1157717230084826e-07, 6.489309513799421e-07, 3.9678959099167477e-07]\n",
      "[0.7785285313924154, 0.7866193453470866, 3.2779879191092715e-07, 1.5803973707641966e-07, 9.17336270770009e-07, 5.609677824161414e-07, 4.6358560799820115e-07, 2.2350584995553053e-07, 1.2973320930947003e-06, 7.933421260120288e-07]\n",
      "[3.6184581915537515, 3.6159478028615317, 1.638050938842924e-06, 7.910830658404983e-07, 4.58517723554586e-06, 2.80502061379283e-06, 2.3165266312477897e-06, 1.1187472538815787e-06, 6.484343632583235e-06, 3.966851578889276e-06]\n",
      "[7.821764230728149, 7.412562131881714, 3.2755421906797827e-06, 1.582071755320782e-06, 9.169010640324028e-06, 5.609276825476654e-06, 4.632448897708904e-06, 2.23745143014359e-06, 1.2967310680567147e-05, 7.932942619716568e-06]\n"
     ]
    }
   ],
   "source": [
    "conv_coeffs = [1e-3, 1e-2, 1e-1, 1, 10, 40, 100, 500, 1000, 5000, 10000] # Values of convection coefficient\n",
    "n_c = 400 # Number of data points for convection coefficient 100\n",
    "seeds = [1, 2, 3]\n",
    "experiments = []\n",
    "\n",
    "n_IBF = 380\n",
    "n_OBF = 14\n",
    "\n",
    "# Parameter sampler for ELM: Sample weights from a normal distribution and biases uniformly from [-4, 4]\n",
    "def sample_parameters_randomly(x, _, rng):\n",
    "    weights = rng.normal(loc=0, scale=1, size=(x.shape[1], n_IBF))\n",
    "    biases = rng.uniform(low=-4, high=4, size=(1, n_IBF)) # low=-np.pi, high=np.pi,  2 * np.pi\n",
    "    idx0 = None\n",
    "    idx1 = None\n",
    "    return weights, biases, idx0, idx1\n",
    "\n",
    "# Periodic basis functions\n",
    "def periodic_outer_basis_sine(x_space, n_outer_basis=None, initial_condition=None):\n",
    "    return np.column_stack([\n",
    "        sine_cos(x_space * (i+1))\n",
    "        for i in range(n_outer_basis//2)\n",
    "        for sine_cos in [np.sin, np.cos]\n",
    "    ])\n",
    "\n",
    "for conv_coeff in conv_coeffs:\n",
    "    # Ground truth for the selected convection coefficient\n",
    "    u_true_test = analytical_sol(x_space_test, t_eval, conv_coeff).T\n",
    "\n",
    "    # Loop over different seeds\n",
    "    rmse_swim = np.ones((len(seeds), ))\n",
    "    rmse_elm = np.ones((len(seeds), ))\n",
    "    rel_err_swim = np.ones((len(seeds), ))\n",
    "    rel_err_elm = np.ones((len(seeds), ))\n",
    "    time_swim = np.ones((len(seeds), ))\n",
    "    time_elm = np.ones((len(seeds), ))\n",
    "    j = 0\n",
    "    info = []\n",
    "    for seed in seeds:\n",
    "        ansatz_swim = BoundaryCompliantAnsatz(\n",
    "                n_outer_basis=n_OBF,\n",
    "                n_inner_basis=n_IBF,\n",
    "                activation=\"tanh\",\n",
    "                random_state=seed,\n",
    "                regularization_scale=1e-10,\n",
    "                target_gen=periodic_outer_basis_sine,\n",
    "            )\n",
    "        ansatz_elm = BoundaryCompliantAnsatz(\n",
    "                n_outer_basis=n_OBF,\n",
    "                n_inner_basis=n_IBF,\n",
    "                activation=\"tanh\",\n",
    "                random_state=seed,\n",
    "                regularization_scale=1e-10,\n",
    "                target_gen=periodic_outer_basis_sine,\n",
    "                parameter_sampler = sample_parameters_randomly\n",
    "        )\n",
    "        \n",
    "        # Collocation points in space\n",
    "        x_space = np.linspace(x_lim[0], x_lim[1], n_c).reshape((-1, 1))\n",
    "\n",
    "        # interior points\n",
    "        x_space_inner = x_space[1:-1]\n",
    "        interior_points = x_space_inner\n",
    "\n",
    "        # Domain\n",
    "        domain = Domain(\n",
    "            interior_points=interior_points,\n",
    "            boundary_points=boundary_points,\n",
    "        )\n",
    "        \n",
    "        adv_solver_swim = AdvectionSolver(\n",
    "            domain=domain, \n",
    "            ansatz=ansatz_swim,\n",
    "            u0=u0,\n",
    "            boundary_condition=boundary_condition,\n",
    "            c=conv_coeff,\n",
    "            forcing=forcing,\n",
    "            regularization_scale=1e-10,\n",
    "        )\n",
    "        adv_solver_elm = AdvectionSolver(\n",
    "            domain=domain, \n",
    "            ansatz=ansatz_elm,\n",
    "            u0=u0,\n",
    "            boundary_condition=boundary_condition,\n",
    "            c=conv_coeff,\n",
    "            forcing=forcing,\n",
    "            regularization_scale=1e-10,\n",
    "        )\n",
    "        \n",
    "        # Compute weights and biases of the SWIM network\n",
    "        time_blocks = 1\n",
    "        t_swim_start = time.time()\n",
    "        sol_swim, solver_status_swim = adv_solver_swim.fit(t_span=[0, np.max(t_eval)], rtol = 1e-10, atol = 1e-10, svd_cutoff=1e-12);\n",
    "        t_swim_stop = time.time()\n",
    "        time_swim[j] = t_swim_stop - t_swim_start\n",
    "        \n",
    "        t_elm_start = time.time()\n",
    "        sol_elm, solver_status_elm = adv_solver_elm.fit(t_span=[0, np.max(t_eval)], rtol = 1e-10, atol = 1e-10, svd_cutoff=1e-12);\n",
    "        t_elm_stop = time.time()\n",
    "        time_elm[j] = t_elm_stop - t_elm_start\n",
    "        \n",
    "        # Evaluate on test data\n",
    "        u_elm_test = adv_solver_elm.evaluate(x_eval=x_space_test, t_eval = t_eval, solver_status=solver_status_elm).T\n",
    "        u_swim_test = adv_solver_swim.evaluate(x_eval=x_space_test, t_eval = t_eval, solver_status=solver_status_swim).T #, solver_status=solver_status\n",
    "                    \n",
    "        # Compute metrics\n",
    "        rmse_elm[j] = np.sqrt(mean_squared_error(u_true_test, u_elm_test))\n",
    "        rmse_swim[j] = np.sqrt(mean_squared_error(u_true_test, u_swim_test))  # mean squared error\n",
    "        rel_err_elm[j] = rmse_elm[j]/np.sqrt(mean_squared_error(u_true_test, np.zeros_like(u_true_test)))\n",
    "        rel_err_swim[j] = rmse_swim[j]/np.sqrt(mean_squared_error(u_true_test, np.zeros_like(u_true_test)))\n",
    "        j += 1\n",
    "\n",
    "    info.append(np.mean(time_elm))\n",
    "    info.append(np.mean(time_swim))\n",
    "    info.append(np.mean(rmse_elm))\n",
    "    info.append(np.std(rmse_elm))\n",
    "    info.append(np.mean(rmse_swim))\n",
    "    info.append(np.std(rmse_swim))\n",
    "    info.append(np.mean(rel_err_elm))\n",
    "    info.append(np.std(rel_err_elm))\n",
    "    info.append(np.mean(rel_err_swim))\n",
    "    info.append(np.std(rel_err_swim))\n",
    "\n",
    "    print(info)\n",
    "    experiments.append(info) \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(11, 10)\n",
      "Information for beta = 40:\n",
      "time_swim =  3.6159478028615317 time_elm =  3.6184581915537515\n",
      "rmse elm =  1.638050938842924e-06 +- 7.910830658404983e-07\n",
      "rmse swim =  4.58517723554586e-06 +- 2.80502061379283e-06\n",
      "rel error elm =  2.3165266312477897e-06 +- 1.1187472538815787e-06\n",
      "rel error swim =  6.484343632583235e-06 +- 3.966851578889276e-06\n"
     ]
    }
   ],
   "source": [
    "res = np.vstack(experiments)\n",
    "print(np.shape(res))\n",
    "\n",
    "# For beta = 40 (Values in the table in the paper)\n",
    "print('Information for beta = 40:')\n",
    "print('time_swim = ', res[-2, -9], 'time_elm = ', res[-2, -10])\n",
    "print('rmse elm = ', res[-2, -8], '+-', res[-2, -7])\n",
    "print('rmse swim = ', res[-2, -6], '+-', res[-2, -5])\n",
    "print('rel error elm = ', res[-2, -4], '+-', res[-2, -3])\n",
    "print('rel error swim = ', res[-2, -2], '+-', res[-2, -1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error Vs flow velocity for SWIM-ODE and ELM-ODE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(11, 10)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27028 characters]",
      "text/plain": [
       "<Figure size 400x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fontsize = 12\n",
    "\n",
    "# Extract errors\n",
    "conv_coeffs = np.reshape(np.asarray(conv_coeffs), (-1, 1))\n",
    "rmse_elm = np.reshape(res[:, -4], (-1, 1))\n",
    "rmse_swim = np.reshape(res[:, -2], (-1, 1))\n",
    "rel_l2_elm = np.reshape(res[:, -8], (-1, 1))\n",
    "rel_l2_swim = np.reshape(res[:, -6], (-1, 1))\n",
    "\n",
    "# Store the errors for different values of beta\n",
    "rmse = np.hstack((conv_coeffs, rmse_elm, rmse_swim))\n",
    "rel_l2 = np.hstack((conv_coeffs, rel_l2_elm, rel_l2_swim))\n",
    "print(np.shape(res))\n",
    "\n",
    "# Store rel l2 error values for convergence plots\n",
    "with open('adv_ode_swim_elm_rmse.npy', 'wb') as f:\n",
    "    np.save(f, rmse)\n",
    "\n",
    "with open('adv_ode_swim_elm_rel_l2.npy', 'wb') as f:\n",
    "    np.save(f, rel_l2)\n",
    "\n",
    "# Visualiza errors vs convection coefficient\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4, 3))\n",
    "ax.loglog(conv_coeffs, res[:, -8], label='rel_err_elm')\n",
    "ax.loglog(conv_coeffs, res[:, -6], label='rel_err_swim')\n",
    "plt.xlabel('beta')\n",
    "plt.ylabel('Relative l_2 error')\n",
    "plt.tick_params(axis='both', labelsize=fontsize)\n",
    "ax.legend()\n",
    "fig.tight_layout()\n",
    "plt.savefig('advection_ode_swim_elm.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
}
