{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Long term simulation with SWIM-ODE and ELM-ODE for travelling sine-wave (t = 100)"
   ]
  },
  {
   "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 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",
    "import time\n",
    "import matplotlib.cm as cm\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_space = 256 # Number of test points in space\n",
    "n_t = 100 #Number of test points in time\n",
    "t_eval = np.linspace(0, 1000, n_t) # Time domain\n",
    "x_lim = [0, 2 * np.pi] # Space domain: 1-d space with uniform grid\n",
    "n_s = 1000 # Number of collocation points in space\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",
    "# Analytical solution\n",
    "def analytical_sol(x, t, conv_coeff):\n",
    "    return np.sin(x - conv_coeff * t)\n",
    "\n",
    "x_space = np.linspace(x_lim[0], x_lim[1], n_s).reshape((-1, 1))\n",
    "# interior points\n",
    "x_space_inner = x_space[1:-1]\n",
    "interior_points = x_space_inner\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",
    "domain = Domain(\n",
    "    interior_points=interior_points,\n",
    "    boundary_points=boundary_points,\n",
    ")\n",
    "\n",
    "# Set boundary conditions\n",
    "boundary_condition = \"periodic\"\n",
    "\n",
    "# Outer 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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyper-parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SWIM and ELM: Hyper-parameters (First entry: SWIM, second entry: ELM)\n",
    "reg_OBF = [1e-10, 1e-10] \n",
    "svd_cutoff = [1e-10, 1e-10] \n",
    "reg_pinv_ic = reg_OBF # Regularization const for pseudo-inverse and least squares solution for IC\n",
    "conv_coeff = 1. # beta (change convection coefficient!!)\n",
    "n_IBFs_factors = [10, 10]\n",
    "ratios = [40, 40]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit and compute errors (averaged over 3 seeds) for SWIM-ODE and ELM-ODE and "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time ELM:  1.4471526145935059 time_swim 1.4797428448994954\n"
     ]
    }
   ],
   "source": [
    "# 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_elm))\n",
    "    biases = rng.uniform(low=-4, high=4, size=(1, n_IBF_elm)) # low=-np.pi, high=np.pi,  2 * np.pi\n",
    "    idx0 = None\n",
    "    idx1 = None\n",
    "    return weights, biases, idx0, idx1\n",
    "\n",
    "# Test Data\n",
    "x_space_test = np.linspace(np.min(x_space), np.max(x_space), n_space).reshape((-1, 1))\n",
    "u_true_test = analytical_sol(x_space_test, t_eval, conv_coeff=conv_coeff).T\n",
    "\n",
    "seeds = [1, 2, 3]\n",
    "# Hyper-parameters for SWIM and ELM\n",
    "n_OBF_swim = int(n_s//ratios[0]) #2 # Number of outer basis functions\n",
    "n_IBF_swim = int(n_IBFs_factors[0] * n_OBF_swim) # Number of inner basis functions\n",
    "n_OBF_elm = int(n_s//ratios[1]) #2 # Number of outer basis functions\n",
    "n_IBF_elm = int(n_IBFs_factors[1] * n_OBF_elm) # Number of inner basis functions\n",
    "\n",
    "# Loop over different seeds\n",
    "rmse_swim = np.ones((len(seeds), ))\n",
    "rmse_swim_rs = np.ones((len(seeds), ))\n",
    "rmse_elm = np.ones((len(seeds), ))\n",
    "rel_err_swim = np.ones((len(seeds), ))\n",
    "rel_err_swim_rs = np.ones((len(seeds), ))\n",
    "rel_err_elm = np.ones((len(seeds), ))\n",
    "time_swim = np.ones((len(seeds), ))\n",
    "time_swim_rs = np.ones((len(seeds), ))\n",
    "time_elm = np.ones((len(seeds), ))\n",
    "l2_rel_err_elm = np.ones((len(seeds), ))\n",
    "l2_rel_err_swim = np.ones((len(seeds), ))\n",
    "\n",
    "j = 0\n",
    "rmse_elms = []\n",
    "rmse_swims = []\n",
    "l2_rel_err_swims = []\n",
    "l2_rel_err_elms = []\n",
    "\n",
    "for s in seeds:\n",
    "    # Set inner basis functions\n",
    "    ansatz_elm = BoundaryCompliantAnsatz(\n",
    "        n_outer_basis=n_OBF_elm,\n",
    "        n_inner_basis=n_IBF_elm,\n",
    "        activation=\"tanh\",\n",
    "        target_gen=periodic_outer_basis_sine,\n",
    "        random_state=s,\n",
    "        regularization_scale=reg_OBF[1],\n",
    "        parameter_sampler = sample_parameters_randomly\n",
    "    )\n",
    "    ansatz_swim = BoundaryCompliantAnsatz(\n",
    "        n_outer_basis=n_OBF_swim,\n",
    "        n_inner_basis=n_IBF_swim,\n",
    "        activation=\"tanh\",\n",
    "        target_gen=periodic_outer_basis_sine,\n",
    "        random_state=s,\n",
    "        regularization_scale=reg_OBF[0],\n",
    "    )\n",
    "\n",
    "    # Set advection solver\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=reg_OBF[1],\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=reg_OBF[0],\n",
    "    )\n",
    "\n",
    "    # Compute weights and biases of the SWIM network\n",
    "    time_blocks = 1\n",
    "    #sol, solver_status = adv_solver_swim.fit_time_blocks(t_span=[0, np.max(t_eval)], time_blocks=time_blocks,rtol = 1e-10, atol = 1e-10);\n",
    "    t_swim_start = time.time()\n",
    "    sol, solver_status = adv_solver_swim.fit(t_span=[0, np.max(t_eval)], rtol = 1e-8, atol = 1e-8, svd_cutoff=1e-11);\n",
    "    t_swim_stop = time.time()\n",
    "    time_swim[j] = t_swim_stop - t_swim_start\n",
    "    \n",
    "    t_elm_start = time.time()\n",
    "    adv_solver_elm.fit(t_span=[0, np.max(t_eval)], rtol = 1e-8, atol = 1e-8, svd_cutoff=1e-11);\n",
    "    t_elm_stop = time.time()\n",
    "    time_elm[j] = t_elm_stop - t_elm_start\n",
    "\n",
    "    # Evaluate solution\n",
    "    u_elm_test = adv_solver_elm.evaluate(x_eval=x_space_test, t_eval = t_eval).T\n",
    "    u_swim_test = adv_solver_swim.evaluate(x_eval=x_space_test, t_eval = t_eval).T #, solver_status=solver_status\n",
    "\n",
    "    # Store useful info for plotting\n",
    "    mse_elm = mean_squared_error(u_true_test, u_elm_test)  # mean squared error\n",
    "    rmse_elm[j] = np.sqrt(mse_elm)  # Root Mean Squared Error\n",
    "    mse_swim = mean_squared_error(u_true_test, u_swim_test)  # mean squared error\n",
    "    rmse_swim[j] = np.sqrt(mse_swim)  # Root Mean Squared Error\n",
    "    l2_rel_err_elm[j] = np.linalg.norm(u_elm_test - u_true_test, 2)/np.linalg.norm(u_true_test, 2)\n",
    "    l2_rel_err_swim[j] = np.linalg.norm(u_swim_test - u_true_test, 2)/np.linalg.norm(u_true_test, 2)\n",
    "    j += 1\n",
    "\n",
    "    # All metrics\n",
    "    rmse_elm = np.ones((np.shape(t_eval))) * 100\n",
    "    rmse_swim = np.ones((np.shape(t_eval))) * 100\n",
    "    l2_rel_err_swim = np.ones((np.shape(t_eval))) * 100\n",
    "    l2_rel_err_elm = np.ones((np.shape(t_eval))) * 100\n",
    "    for t in range(len(t_eval)):\n",
    "        mse_elm = mean_squared_error(u_true_test[t, :], u_elm_test[t, :])  # mean squared error\n",
    "        rmse_elm[t] = np.sqrt(mse_elm)  # Root Mean Squared Error\n",
    "        mse_swim = mean_squared_error(u_true_test[t, :], u_swim_test[t, :])  # mean squared error\n",
    "        rmse_swim[t] = np.sqrt(mse_swim)  # Root Mean Squared Error\n",
    "        l2_rel_err_elm[t] = np.linalg.norm(u_elm_test[t, :] - u_true_test[t, :], 2)/np.linalg.norm(u_true_test[t, :], 2)\n",
    "        l2_rel_err_swim[t] = np.linalg.norm(u_swim_test[t, :] - u_true_test[t, :], 2)/np.linalg.norm(u_true_test[t, :], 2)\n",
    "    rmse_elms.append(rmse_elm)\n",
    "    rmse_swims.append(rmse_swim)\n",
    "    l2_rel_err_elms.append(l2_rel_err_elm)\n",
    "    l2_rel_err_swims.append(l2_rel_err_swim)\n",
    "\n",
    "time_elm = np.mean(time_elm)\n",
    "time_swim = np.mean(time_swim)\n",
    "\n",
    "print('Time ELM: ', time_elm, 'time_swim', time_swim)\n",
    "\n",
    "rmse_elm_avg = (rmse_elms[0] + rmse_elms[1] + rmse_elms[2]) / 3\n",
    "rmse_swim_avg = (rmse_swims[0] + rmse_swims[1] + rmse_swims[2])/3\n",
    "l2_rel_err_elms_avg =  (l2_rel_err_elms[0] +  l2_rel_err_elms[1] +  l2_rel_err_elms[2])/3\n",
    "l2_rel_err_swims_avg = (l2_rel_err_swims[0] + l2_rel_err_swims[1] + l2_rel_err_swims[2])/3\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error plot for Long-term simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36848 characters]",
      "text/plain": [
       "<Figure size 400x250 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 79644 characters]",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot of error vs time for SWIM-ODE and ELM-ODE\n",
    "# Figure parameters\n",
    "from matplotlib.ticker import LogLocator, MaxNLocator\n",
    "rc_fonts = {\n",
    "    \"text.usetex\": False,\n",
    "}\n",
    "plt.rcParams.update(rc_fonts)\n",
    "fontsize = 14\n",
    "cmap =cm.jet\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4, 2.5), sharey=True)\n",
    "ax.semilogy(t_eval, l2_rel_err_elms_avg, label='ELM-ODE', c='green')\n",
    "ax.semilogy(t_eval, l2_rel_err_swims_avg, label='SWIM-ODE', c='red')\n",
    "ax.legend(frameon=False, loc='lower right', fontsize=fontsize)\n",
    "ax.set_xlabel('Time', fontsize=fontsize)\n",
    "plt.ylabel(r'Relative  ' + r'$\\mathbb{L}_{2}$ ' + r'error', fontsize=fontsize)\n",
    "ax.yaxis.set_major_locator(LogLocator(base=10.0, numticks=5))  # Major ticks at log scale\n",
    "ax.yaxis.set_minor_locator(LogLocator(base=10.0, subs='auto', numticks=10))  # Minor ticks\n",
    "ax.xaxis.set_major_locator(MaxNLocator(nbins=5))\n",
    "ax.tick_params(axis='both', which='major', labelsize=14)\n",
    "ax.tick_params(axis='both', which='minor', labelsize=10)\n",
    "fig.tight_layout()\n",
    "plt.savefig('adv_long_time_simulation.pdf')\n",
    "\n",
    "\n",
    "# visualize the solution\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6, 3), constrained_layout=True)#\n",
    "extent = [np.min(t_eval) , np.max(t_eval), np.min(x_space), np.max(x_space)]\n",
    "aspect = 16\n",
    "sol_img1 = ax.imshow(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=50)#, fraction=0.046 * 3 \n",
    "ax.set_title('True solution', fontsize=fontsize)\n",
    "cb.ax.tick_params(labelsize=fontsize)\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_true_sol.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
}
