{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import reservoirpy as rpy\n",
    "import reservoirpy.nodes as rpn\n",
    "from reservoirpy.datasets import lorenz, rossler, doublescroll, kuramoto_sivashinsky\n",
    "\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "rpy.verbosity(False)\n",
    "\n",
    "rpy.set_seed(42)\n",
    "\n",
    "name_idx = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_noise(time_series, noise_type=\"random\", noise_level=0.5):\n",
    "    \"\"\"\n",
    "    对输入的时间序列添加噪声。\n",
    "    \n",
    "    参数:\n",
    "        time_series (numpy.ndarray): 输入的时间序列。\n",
    "        noise_type (str): 噪声类型，可选值为 \"random\"（随机噪声）, \"sin\"（正弦噪声）, 或 \"gaussian\"（正态分布噪声）。\n",
    "        noise_level (float): 噪声强度，决定噪声幅度。\n",
    "    \n",
    "    返回:\n",
    "        numpy.ndarray: 添加噪声后的时间序列。\n",
    "    \"\"\"\n",
    "    if noise_type == \"random\":\n",
    "        noise = np.random.uniform(-noise_level, noise_level, len(time_series))\n",
    "    elif noise_type == \"sin\":\n",
    "        noise = noise_level * np.sin(30 * np.pi * np.arange(len(time_series)) / len(time_series))\n",
    "    elif noise_type == \"gaussian\":\n",
    "        noise = np.random.normal(0, noise_level, len(time_series))\n",
    "    else:\n",
    "        raise ValueError(\"Unsupported noise_type. Choose from 'random', 'sin', or 'gaussian'.\")\n",
    "    \n",
    "    return time_series + np.stack((noise,noise,noise)).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show(t1, t2, figsize=(18, 6), title=\"Lorenz Attractor\"):\n",
    "    '''\n",
    "    t1: Noisy data(plot with solid line & orange color)\n",
    "    \n",
    "    t2: Clean data(plot with dashed line & blue color)\n",
    "    '''\n",
    "    plt.figure(figsize=figsize)\n",
    "\n",
    "    plt.subplot(3, 1, 1)\n",
    "    #plt.plot(t2[:,0])\n",
    "    plt.plot(t1[:,0], color='#ff7f0e')\n",
    "    plt.plot(t2[:,0], linestyle='--')\n",
    "    plt.title('X-component')\n",
    "\n",
    "    plt.subplot(3, 1, 2)\n",
    "    #plt.plot(t2[:,1])\n",
    "    plt.plot(t1[:,1], color='#ff7f0e')\n",
    "    plt.plot(t2[:,1], linestyle='--')\n",
    "    plt.title('Y-component')\n",
    "\n",
    "    plt.subplot(3, 1, 3)\n",
    "    #plt.plot(t2[:,2])\n",
    "    plt.plot(t1[:,2], color='#ff7f0e')\n",
    "    plt.plot(t2[:,2], linestyle='--')\n",
    "    plt.title('Z-component')\n",
    "    \n",
    "    fig = plt.figure(figsize=(12,5), dpi=150)\n",
    "    \n",
    "    rmse = np.sqrt(np.mean((t1 - t2)**2))\n",
    "    fig.suptitle(f\"{title}\\nRMSE: {rmse:.8f}\")\n",
    "    \n",
    "    ax = fig.add_subplot(121, projection='3d')\n",
    "    ax.plot(t2[:,0], t2[:,1], t2[:,2])\n",
    "    ax.set_title(\"Raw Data\")\n",
    "\n",
    "    ax = fig.add_subplot(122, projection='3d')\n",
    "    ax.plot(t1[:,0], t1[:,1], t1[:,2], color='#ff7f0e')\n",
    "    ax.set_title(\"Processed Data\")\n",
    "    \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data acquisition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm01 = lambda x: (x - np.min(x, axis=0)) / (np.max(x, axis=0) - np.min(x, axis=0))\n",
    "\n",
    "lorenz_data = norm01(lorenz(12000, h=0.01)[2000:,:])\n",
    "rossler_data = norm01(rossler(12000, h=0.02)[2000:,:])\n",
    "doublescroll_data = norm01(doublescroll(12000, h=0.1)[2000:,:])\n",
    "\n",
    "noisy_lorenz_data = add_noise(lorenz_data, noise_type=\"gaussian\", noise_level=0.8)\n",
    "noisy_rossler_data = add_noise(rossler_data, noise_type=\"gaussian\", noise_level=0.8)\n",
    "noisy_doublescroll_data = add_noise(doublescroll_data, noise_type=\"gaussian\", noise_level=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(noisy_lorenz_data, lorenz_data, title=\"Lorenz Raw Data & Noisy Data\")\n",
    "show(noisy_rossler_data, rossler_data, title=\"Rossler Raw Data & Noisy Data\")\n",
    "show(noisy_doublescroll_data, doublescroll_data, title=\"Doublescroll Raw Data & Noisy Data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lorenz Test Case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lorenz_train_input = noisy_lorenz_data[:7000]\n",
    "lorenz_train_target = lorenz_data[:7000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lorenz_res = rpn.Reservoir(units=1000, lr=0.5, sr=0.9, activation='tanh', equation='external')\n",
    "lorenz_readout = rpn.Ridge(ridge=5e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train readout\n",
    "states = lorenz_res.run(lorenz_train_input)\n",
    "lorenz_readout.fit(states, lorenz_train_target)\n",
    "output = lorenz_readout.run(states)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show the Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'RMSE of output: {np.sqrt(np.mean((output - lorenz_train_target)**2))}')\n",
    "show(output, lorenz_train_target, title=\"Lorenz RC Output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lorenz_test_input = noisy_lorenz_data[7000:]\n",
    "lorenz_test_target = lorenz_data[7000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "states = lorenz_res.run(lorenz_test_input)\n",
    "output = lorenz_readout1.run(states)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show the Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'RMSE of output: {np.sqrt(np.mean((output - lorenz_test_target)**2))}')\n",
    "show(output1, lorenz_test_target, title=\"Lorenz RC Output\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 迁移学习"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_lorenz_data = norm01(lorenz(12000, h=0.01, x0=[0.1, 0.2, 0.3])[2000:,:])\n",
    "noisy_new_lorenz_data = add_noise(new_lorenz_data, noise_type=\"gaussian\", noise_level=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_lorenz_test_input = noisy_new_lorenz_data\n",
    "new_lorenz_test_target = new_lorenz_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = lorenz_res.run(new_lorenz_test_input)\n",
    "output = lorenz_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(output, new_lorenz_test_target, title=\"Lorenz RC Output\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rossler Test Case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rossler_train_input = noisy_rossler_data[:7000]\n",
    "rossler_train_target = rossler_data[:7000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rossler_input = rpn.Input()\n",
    "rossler_res = rpn.Reservoir(units=1000, lr=0.5, sr=0.9, activation='tanh', equation='external')\n",
    "rossler_readout = rpn.Ridge(ridge=5e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = rossler_res.run(rossler_train_input)\n",
    "rossler_readou1.fit(states, rossler_train_target)\n",
    "output = rossler_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(rossler_train_input, rossler_train_target, title=\"Rossler Raw Data & Noisy Data\")\n",
    "show(output, rossler_train_target, title=\"Rossler RC Output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rossler_test_input = noisy_rossler_data[7000:]\n",
    "rossler_test_target = rossler_data[7000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = rossler_res.run(rossler_test_input)\n",
    "output = rossler_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(output, rossler_test_target, title=\"Rossler RC Output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_rossler_data = norm01(rossler(12000, h=0.02, x0=[0.1, 0.2, 0.3])[2000:,:])\n",
    "noisy_new_rossler_data = add_noise(new_rossler_data, noise_type=\"gaussian\", noise_level=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_rossler_test_input = noisy_new_rossler_data\n",
    "new_rossler_test_target = new_rossler_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = rossler_res.run(new_rossler_test_input)\n",
    "output = rossler_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(output, new_rossler_test_target, title=\"Rossler RC Output\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Double-scroll test case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doublescroll_train_input = noisy_doublescroll_data[:7000]\n",
    "doublescroll_train_target = doublescroll_data[:7000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doublescroll_input = rpn.Input()\n",
    "doublescroll_res = rpn.Reservoir(units=1000, lr=0.5, sr=0.9, activation='tanh', equation='external')\n",
    "doublescroll_readout = rpn.Ridge(ridge=5e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = doublescroll_res.run(doublescroll_train_input)\n",
    "doublescroll_readout1.fit(states, doublescroll_train_target)\n",
    "output = doublescroll_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(doublescroll_train_input, doublescroll_train_target, title=\"Doublescroll Raw Data & Noisy Data\")\n",
    "show(output, doublescroll_train_target, title=\"Doublescroll RC Output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doublescroll_test_input = noisy_doublescroll_data[7000:]\n",
    "doublescroll_test_target = doublescroll_data[7000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = doublescroll_res.run(doublescroll_test_input)\n",
    "output = doublescroll_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(output, doublescroll_test_target, title=\"Doublescroll RC Output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_doublescroll_data = norm01(doublescroll(12000, h=0.1, x0=[0.1, 0.2, 0.3])[2000:,:])\n",
    "noisy_new_doublescroll_data = add_noise(new_doublescroll_data, noise_type=\"gaussian\", noise_level=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_doublescroll_test_input = noisy_new_doublescroll_data\n",
    "new_doublescroll_test_target = new_doublescroll_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = doublescroll_res.run(new_doublescroll_test_input)\n",
    "output = doublescroll_readout.run(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(output, new_doublescroll_test_target, title=\"Doublescroll RC Output\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}