{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Euler Bernoulli Beam Equation: ODE-ELM, ODE-SWIM solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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 EulerBernoulliSolver\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import mean_squared_error\n",
    "\n",
    "#testing\n",
    "import time\n",
    "import matplotlib.cm as cm\n",
    "cmap = cm.jet\n",
    "import matplotlib.ticker as ticker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Define the domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of test points in space\n",
    "n_test = 256\n",
    "t_eval = np.linspace(0, 1, 100)\n",
    "# Set time domain\n",
    "n_t = 100\n",
    "\n",
    "# Domain: 1-d space with uniform grid\n",
    "x_lim = [0, np.pi]\n",
    "\n",
    "# initial condition\n",
    "def u0(x):\n",
    "    return np.sin(x)\n",
    "\n",
    "def ut0(x):\n",
    "    return np.zeros_like(x)\n",
    "\n",
    "# boundary condition\n",
    "boundary_condition = \"periodic\" #\"zero dirichlet second order\"  #\"periodic strict\" #strict\n",
    "\n",
    "# forcing\n",
    "def forcing(x, t):\n",
    "    return (1. - 16* np.pi**2 ) * np.sin(x) * np.cos(4 * np.pi * t)\n",
    "\n",
    "# Analytical solution\n",
    "def analytical_sol(x, t):\n",
    "    return np.sin(x) * np.cos(4 * np.pi * t)\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(x_space * (i+1))\n",
    "        for i in range(n_outer_basis)\n",
    "        for sine in [np.sin]\n",
    "    ])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit the SWIM network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Integration successful.\n",
      "Integration successful.\n",
      "Integration successful.\n",
      "Integration successful.\n",
      "Integration successful.\n",
      "Integration successful.\n",
      "elm time =  0.08053151766459148 swim time 0.0851132074991862\n",
      "rmse elm =  1.7576782931737323e-08 +- 3.910091870412992e-09\n",
      "rmse swim =  6.062669147553227e-08 +- 2.966283572970679e-08\n",
      "rel l-2 error elm 3.504762477371116e-08 +- 7.79661632263372e-09\n",
      "rel l-2 error swim 1.2088796581024773e-07 +- 5.9146883728183974e-08\n"
     ]
    }
   ],
   "source": [
    "experiments = []\n",
    "# List of all hyper-parameters\n",
    "n_x = 400\n",
    "reg_pinv_ic =1e-12\n",
    "reg_pinv_ic = 1e-12 \n",
    "reg_OBF = 1e-12\n",
    "# Data points / number of basis functions\n",
    "seeds = [1, 2, 3]\n",
    "n_IBF = 200\n",
    "n_OBF = 10\n",
    "experiments = []\n",
    "\n",
    "x_space = np.linspace(x_lim[0], x_lim[1], n_x).reshape((-1, 1))\n",
    "\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",
    "# 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=-2, high=2, 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",
    "# Set boundary conditions\n",
    "boundary_condition = \"periodic\"\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",
    "    # Set inner basis functions\n",
    "    ansatz_elm = BoundaryCompliantAnsatz(\n",
    "        n_outer_basis=n_OBF,\n",
    "        n_inner_basis=n_IBF,\n",
    "        activation=\"tanh\",\n",
    "        target_gen=periodic_outer_basis_sine,\n",
    "        random_state=seed,\n",
    "        regularization_scale=reg_OBF,\n",
    "        parameter_sampler=sample_parameters_randomly,\n",
    "        domain_margin_percent=0.1,\n",
    "        scale_bc = 100.,\n",
    "\n",
    "    )\n",
    "    ansatz_swim = BoundaryCompliantAnsatz(\n",
    "        n_outer_basis=n_OBF,\n",
    "        n_inner_basis=n_IBF,\n",
    "        activation=\"tanh\",\n",
    "        target_gen=periodic_outer_basis_sine,\n",
    "        random_state=seed,\n",
    "        regularization_scale=reg_OBF,\n",
    "    )\n",
    "\n",
    "    # Set advection solver\n",
    "    euler_bernoulli_solver_swim = EulerBernoulliSolver(\n",
    "        domain=domain, \n",
    "        ansatz=ansatz_swim,\n",
    "        u0=u0,\n",
    "        ut0=ut0,\n",
    "        boundary_condition=boundary_condition,\n",
    "        forcing=forcing,\n",
    "        regularization_scale=reg_pinv_ic\n",
    "    )\n",
    "    euler_bernoulli_solver_elm = EulerBernoulliSolver(\n",
    "        domain=domain, \n",
    "        ansatz=ansatz_elm,\n",
    "        u0=u0,\n",
    "        ut0=ut0,\n",
    "        boundary_condition=boundary_condition,\n",
    "        forcing=forcing,\n",
    "        regularization_scale=reg_pinv_ic\n",
    "    )\n",
    "    \n",
    "    # Compute weights and biases of the SWIM network\n",
    "    time_blocks = 1\n",
    "    t_swim_start = time.time()\n",
    "    sol= euler_bernoulli_solver_swim.fit(t_span=[0, np.max(t_eval)], rtol = 1e-10, atol = 1e-10); #, num_svd=50\n",
    "    t_swim_stop = time.time()\n",
    "    time_swim[j] = t_swim_stop - t_swim_start\n",
    "    \n",
    "    t_elm_start = time.time()\n",
    "    euler_bernoulli_solver_elm.fit(t_span=[0, np.max(t_eval)], rtol = 1e-10, atol = 1e-10);#, num_svd=50\n",
    "    t_elm_stop = time.time()\n",
    "    time_elm[j] = t_elm_stop - t_elm_start\n",
    "    \n",
    "\n",
    "    # Evaluate solution\n",
    "    x_space_test = np.linspace(np.min(x_space), np.max(x_space), n_test).reshape((-1, 1))\n",
    "    u_true_test = analytical_sol(x_space_test, t_eval).T\n",
    "    u_elm_test = euler_bernoulli_solver_elm.evaluate(x_eval=x_space_test, t_eval = t_eval).T\n",
    "    u_swim_test = euler_bernoulli_solver_swim.evaluate(x_eval=x_space_test, t_eval = t_eval).T\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",
    "    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 = 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",
    "print('elm time = ', np.mean(time_elm), 'swim time', np.mean(time_swim))\n",
    "print('rmse elm = ',np.mean(rmse_elm), '+-', np.std(rmse_elm))\n",
    "print('rmse swim = ',np.mean(rmse_swim), '+-', np.std(rmse_swim))\n",
    "print('rel l-2 error elm',np.mean(rel_err_elm), '+-', np.std(rel_err_elm))\n",
    "print('rel l-2 error swim',np.mean(rel_err_swim), '+-', np.std(rel_err_swim))\n",
    "experiments.append(info)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 76300 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize the solution\n",
    "fig, ax = plt.subplots(1, 3, figsize=(10, 3))\n",
    "plt.rcParams['image.cmap'] = 'jet' #'brg'\n",
    "\n",
    "extent = [np.min(x_space), np.max(x_space), 0 , 1]\n",
    "aspect = 4\n",
    "sol_img1 = ax[0].imshow(u_true_test, extent=extent, origin='lower', aspect=aspect)#, vmin=vlim[0], vmax=vlim[1]\n",
    "sol_img2 = ax[1].imshow(u_swim_test, extent=extent, origin='lower', aspect=aspect)#, vmin=vlim[0], vmax=vlim[1]\n",
    "error_img = ax[2].imshow(abs(u_swim_test - u_true_test),  extent=extent, origin='lower', aspect=aspect)\n",
    "\n",
    "fig.colorbar(sol_img1, ax=ax[0], location='bottom')\n",
    "fig.colorbar(sol_img2, ax=ax[1], location='bottom')\n",
    "fig.colorbar(error_img, ax=ax[2], location='bottom')\n",
    "\n",
    "ax[0].set_title('True solution')\n",
    "ax[1].set_title('Network solution')\n",
    "ax[2].set_title('Difference')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Error plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 82744 characters]",
      "text/plain": [
       "<Figure size 600x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 82744 characters]",
      "text/plain": [
       "<Figure size 600x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fontsize = 24\n",
    "aspect = 0.09\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('bernoulli_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",
    "# Create the colorbar\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",
    "\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('bernoulli_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
}
