{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving non-linear diffusion equation using ODE-ELM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import sys\n",
    "sys.path.append('../../')\n",
    "sys.path.append('../../src')\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from swimpde import Domain\n",
    "from swimpde import BasicAnsatz\n",
    "from swimpde import Nonlinear_Diffusion_Solver\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",
    "from examples import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load point cloud data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read and scale data\n",
    "scaler = MinMaxScaler(feature_range=(0.65, 0.9))\n",
    "\n",
    "# Load and preprocess training data\n",
    "data_boundary = np.load('../../data/nature_inspired_dataset/nature_inspired_boundary_1000.npy')\n",
    "data_interior = np.load('../../data/nature_inspired_dataset/nature_inspired_interior_1000.npy')\n",
    "data_full = np.vstack((data_boundary, data_interior))\n",
    "data_new = scaler.fit_transform(data_full) \n",
    "data_boundary = data_new[:np.shape(data_boundary)[0], :]\n",
    "data_interior = data_new[np.shape(data_boundary)[0]::, :]\n",
    "\n",
    "# Load and preprocess test data\n",
    "data_eval = np.load('../../data/nature_inspired_dataset/nature_inspired_eval.npy')\n",
    "x_eval = scaler.fit_transform(data_eval)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "# initial condition\n",
    "def u0(x):\n",
    "    return np.sin(np.pi * x[:, 0]) * x[:, 1]**-3  \n",
    "\n",
    "# forcing\n",
    "def forcing(x, t):    \n",
    "    return np.exp(-t) * np.sin(np.pi * x[:, 0]) * x[:, 1]**(-3) * (-1. - (-1) * np.exp(-t) * x[:, 1]**(-5) * np.sin(np.pi*x[:, 0]) * (-12 + np.pi**2 * x[:, 1]**2) )\n",
    "\n",
    "# boundary condition\n",
    "boundary_condition = \"dirichlet\"\n",
    "\n",
    "# Analytical solution\n",
    "def analytical_sol(x, t):\n",
    "    return np.exp(-t) * np.sin(np.pi * x[:, 0]) * x[:, 1]**-3  \n",
    "    \n",
    "# Test data\n",
    "t_eval = np.linspace(0, 1, 100).reshape(-1, 1, 1) # time domain\n",
    "x_train = data_interior # space domain\n",
    "u_true =  analytical_sol(x_eval, t_eval)\n",
    "u_true = np.reshape(u_true, (np.shape(u_true)[0], np.shape(u_true)[2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54308 characters]",
      "text/plain": [
       "<Figure size 600x200 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Visualize the true solution\n",
    "timesteps = [0, 30, 60, 99]\n",
    "utils.plot(x=x_eval, u=u_true, timesteps=timesteps, \n",
    "           title='Ground truth', cmap_offset=0,\n",
    "           savefig=True,figname='ground_truth.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit and evaluate ELM-ODE for good hyper-parameters "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time= 7.6800696849823 rmse_elm= 0.009659942956645576 rel_err_elm= 0.00986846563449254\n",
      "time= 14.6868257522583 rmse_elm= 0.007114441694169829 rel_err_elm= 0.007268016351920141\n",
      "time= 10.870229482650757 rmse_elm= 0.005254592481273539 rel_err_elm= 0.0053680198275950705\n",
      "elm-ode time =  11.079041639963785\n",
      "rmse elm-ode =  0.007342992377362981 +- 0.0018057232648454565\n",
      "rel l-2 error elm-ode =  0.007501500604669251 +- 0.0018447021959143092\n"
     ]
    }
   ],
   "source": [
    "seeds = [1, 2, 3]\n",
    "experiments = []\n",
    "\n",
    "def set_hyperparams(case:str):\n",
    "    if case == 'fast':\n",
    "        width = 300 # Width\n",
    "        reg_const = 1e-10 # Regularization constant\n",
    "        svd_cutoff = 1e-10\n",
    "    else:\n",
    "        width = 200 # Width\n",
    "        reg_const = 1e-15 # Regularization constant\n",
    "        svd_cutoff = 1e-15\n",
    "    return width, reg_const, svd_cutoff\n",
    "\n",
    "width, reg_const, svd_cutoff = set_hyperparams('fast')\n",
    "rtol = 1e-6\n",
    "atol = 1e-6\n",
    "# Loop over different seeds\n",
    "rmse_elm = np.ones((len(seeds), ))\n",
    "rel_err_elm = np.ones((len(seeds), ))\n",
    "time_elm = np.ones((len(seeds), ))\n",
    "j = 0\n",
    "info = []\n",
    "for seed in seeds:\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], width))\n",
    "        biases = rng.uniform(low=-4, high=4, size=(1, width)) # low=-np.pi, high=np.pi,  2 * np.pi\n",
    "        idx0 = None\n",
    "        idx1 = None\n",
    "        return weights, biases, idx0, idx1\n",
    "\n",
    "    ansatz_elm = BasicAnsatz(\n",
    "        n_neurons=width,\n",
    "        activation=\"tanh\",\n",
    "        random_state=seed,\n",
    "        regularization_scale=reg_const,\n",
    "        parameter_sampler = sample_parameters_randomly\n",
    "    )\n",
    "    \n",
    "    # Interior points\n",
    "    normal_vectors = data_boundary.copy()\n",
    "    # Domain\n",
    "    domain = Domain(\n",
    "        interior_points=data_interior,\n",
    "        boundary_points=data_boundary,\n",
    "        normal_vectors=normal_vectors\n",
    "    )\n",
    "    nonlinear_diffusion_solver_elm = Nonlinear_Diffusion_Solver(\n",
    "        domain=domain, \n",
    "        ansatz=ansatz_elm,\n",
    "        u0=u0,\n",
    "        boundary_condition=boundary_condition,\n",
    "        forcing=forcing,\n",
    "        regularization_scale=reg_const,\n",
    "        scale_boundary_correction=1.,\n",
    "        boundary_condition_true=analytical_sol\n",
    "    )\n",
    "    # Compute weights and biases of the SWIM network\n",
    "    time_blocks = 1\n",
    "    ic_eval = u0(domain.interior_points)\n",
    "    t_elm_start = time.time()\n",
    "    sol_elm, solver_status_elm = nonlinear_diffusion_solver_elm.fit(t_span=[0, np.max(t_eval)], \n",
    "                                                                       rtol = rtol, atol = atol, svd_cutoff= svd_cutoff,\n",
    "                                                                        outer_basis=False,\n",
    "                                                                        init_cond=ic_eval);\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 = nonlinear_diffusion_solver_elm.evaluate(x_eval=x_eval, t_eval = t_eval).T #, solver_status=solver_status\n",
    "                \n",
    "    # Compute metrics\n",
    "    rmse_elm[j] = np.sqrt(mean_squared_error(u_true, u_elm_test))\n",
    "    rel_err_elm[j] = rmse_elm[j]/np.sqrt(mean_squared_error(u_true, np.zeros_like(u_true)))\n",
    "    info.append(time_elm[j])\n",
    "    info.append(rmse_elm[j])\n",
    "    info.append(rel_err_elm[j])\n",
    "    print('time=', time_elm[j], 'rmse_elm=', rmse_elm[j], 'rel_err_elm=',rel_err_elm[j])\n",
    "    j += 1\n",
    "\n",
    "print('elm-ode time = ', np.mean(time_elm))\n",
    "print('rmse elm-ode = ',np.mean(rmse_elm), '+-', np.std(rmse_elm))\n",
    "print('rel l-2 error elm-ode = ',np.mean(rel_err_elm), '+-', np.std(rel_err_elm))\n",
    "experiments.append(info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55120 characters]",
      "text/plain": [
       "<Figure size 600x200 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the ELM-ODE solution\n",
    "utils.plot(x=x_eval, u=u_elm_test, timesteps=timesteps, title='ELM-ODE solution',cmap_offset=0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 111248 characters]",
      "text/plain": [
       "<Figure size 700x300 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the absolute error with ELM-ODE\n",
    "utils.plot_error(x=x_eval, u_true=u_true, u_nn=u_elm_test, timesteps=timesteps, \n",
    "           figsize=(7,3), fontsize=14, \n",
    "           title='Absolute error: ELM-ODE',savefig=True, \n",
    "           figname='elm_ode_error.png')"
   ]
  }
 ],
 "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
}
