{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving advection equation using ELM-ODE and SWIM-ODE for advection coefficient of 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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",
    "import matplotlib.cm as cm\n",
    "np.random.seed(2)\n",
    "import time\n",
    "cmap = cm.jet\n",
    "import matplotlib.ticker as ticker\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.7152887980143229, 0.7208210627237955, 0.6666666683532934, 0.4714045184057814, 0.6666666811783809, 0.4714045002683687, 0.6666666690518834, 0.47140451741782563, 0.666666687189042, 0.47140449176801014]\n",
      "[0.7152887980143229, 0.7208210627237955, 0.6666666683532934, 0.4714045184057814, 0.6666666811783809, 0.4714045002683687, 0.6666666690518834, 0.47140451741782563, 0.666666687189042, 0.47140449176801014, 0.43882322311401367, 0.4320869445800781, 0.333333340886142, 0.4714045154503895, 0.33333335020188376, 0.47140450886316554, 0.3333333440144681, 0.4714045132383289, 0.3333333571887318, 0.47140450392271815]\n",
      "[0.7152887980143229, 0.7208210627237955, 0.6666666683532934, 0.4714045184057814, 0.6666666811783809, 0.4714045002683687, 0.6666666690518834, 0.47140451741782563, 0.666666687189042, 0.47140449176801014, 0.43882322311401367, 0.4320869445800781, 0.333333340886142, 0.4714045154503895, 0.33333335020188376, 0.47140450886316554, 0.3333333440144681, 0.4714045132383289, 0.3333333571887318, 0.47140450392271815, 0.1534738540649414, 0.15613349278767905, 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"
     ]
    }
   ],
   "source": [
    "conv_coeff = [40] # 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",
    "# 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",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of solutions using SWIM-ODE and ELM-ODE to true solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 135980 characters]",
      "text/plain": [
       "<Figure size 600x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 135980 characters]",
      "text/plain": [
       "<Figure size 600x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fontsize = 22\n",
    "aspect = 0.07\n",
    "\n",
    "# visualize the solution\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6, 5), constrained_layout=True)#\n",
    "extent = [np.min(t_eval) , np.max(t_eval), np.min(x_space), np.max(x_space)]\n",
    "sol_img1 = ax.imshow(abs(u_swim_test-u_true_test).T, extent=extent, origin='lower', aspect=aspect, cmap=cmap)#, vmin=vlim[0], vmax=vlim[1]\n",
    "cb = fig.colorbar(sol_img1, ax=ax, location='bottom',aspect=20)#, fraction=0.046 * 3 \n",
    "\n",
    "# Set the formatter for the colorbar to scientific notation\n",
    "cb.formatter = ticker.FuncFormatter(lambda x, pos: f'{x:.1e}')\n",
    "cb.update_ticks()\n",
    "\n",
    "# Set the font size of the colorbar labels\n",
    "cb.ax.tick_params(labelsize=fontsize)  # Change 12 to your desired font size\n",
    "ax.set_xlabel('t', fontsize=fontsize)\n",
    "ax.set_ylabel('x', fontsize=fontsize)\n",
    "plt.tick_params(axis='both', labelsize=fontsize)\n",
    "plt.savefig('adv_error_swim.pdf')\n",
    "\n",
    "# visualize the solution\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6, 5), constrained_layout=True)#\n",
    "extent = [np.min(t_eval) , np.max(t_eval), np.min(x_space), np.max(x_space)]\n",
    "sol_img1 = ax.imshow(abs(u_swim_test-u_true_test).T, extent=extent, origin='lower', aspect=aspect, cmap=cmap)#, vmin=vlim[0], vmax=vlim[1]\n",
    "cb = fig.colorbar(sol_img1, ax=ax, location='bottom',aspect=20)#, fraction=0.046 * 3 \n",
    "\n",
    "# Set the formatter for the colorbar to scientific notation\n",
    "cb.formatter = ticker.FuncFormatter(lambda x, pos: f'{x:.1e}')\n",
    "cb.update_ticks()\n",
    "\n",
    "# Set the font size of the colorbar labels\n",
    "cb.ax.tick_params(labelsize=fontsize)  # Change 12 to your desired font size\n",
    "ax.set_xlabel('t', fontsize=fontsize)\n",
    "ax.set_ylabel('x', fontsize=fontsize)\n",
    "plt.tick_params(axis='both', labelsize=fontsize)\n",
    "plt.savefig('adv_error_elm.pdf')\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
}
