spiral_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e92b0f1",
   "metadata": {},
   "source": [
    "# import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e9bf230f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "93b75332",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# import packages\n",
    "import os\n",
    "import argparse\n",
    "import logging\n",
    "import time\n",
    "import numpy as np\n",
    "import numpy.random as npr\n",
    "import matplotlib\n",
    "# matplotlib.use('agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "import random\n",
    "\n",
    "from datetime import timedelta\n",
    "import scipy.stats\n",
    "from sklearn.decomposition import PCA\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import copy\n",
    "\n",
    "from scipy.interpolate import splev, splrep, interp1d\n",
    "\n",
    "import sklearn\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "\n",
    "from prettytable import PrettyTable\n",
    "\n",
    "from torchdiffeq import odeint_adjoint as odeint_adjoint\n",
    "from torchdiffeq import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "857465c5",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# import functions and NNs\n",
    "\n",
    "from adaptive_select import create_adap_data_buffer\n",
    "from interpolate_windows import interpolate_window, interpolate_window_adap\n",
    "from LSTM_functions import fit_LSTM_grids, pred_LSTM_grids, train_LSTM\n",
    "from RNN_functions import fit_RNN_grids, pred_RNN_grids, train_RNN\n",
    "from RNN_ODE_functions import fit_non_adap_models_grids, pred_non_adap_models_grids, train_non_adap_models\n",
    "from RNN_ODE_Adap_functions import fit_adap_models, pred_adap_models, train_adap_models\n",
    "from NN import RNNODE, OutputNN, RNN, LSTM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7604b0b6",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def seed_everything(seed=1234):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.use_deterministic_algorithms(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4f68d5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set parameters\n",
    "n_latent = 128\n",
    "n_hidden = 128\n",
    "time_scale = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e0e1af",
   "metadata": {},
   "source": [
    "# Generate spiral data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "857bcbf2",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "class Lambda(nn.Module):\n",
    "    \n",
    "    def __init__(self, true_A, alpha=2, beta = 3, gamma = 0, const = 1.):\n",
    "        super(Lambda, self).__init__()\n",
    "        self.alpha = alpha\n",
    "        self.beta = beta\n",
    "        self.gamma = gamma\n",
    "        self.true_A = true_A\n",
    "        self.const = const\n",
    "        \n",
    "    def orig_f(self, t, y):\n",
    "        return torch.mm(y, self.true_A) * (1/torch.norm(y)**self.alpha) #torch.mm(y**3, self.true_A) * (1/torch.norm(y)**self.alpha) \n",
    "    \n",
    "    def sec_order_deriv(self, t, y):\n",
    "        return (3/torch.norm(y)**self.alpha*self.orig_f(t, y) @ torch.diag((y.reshape((2,)))**2) @ self.true_A\n",
    "              - self.alpha/torch.norm(y)**2 * torch.dot(y.reshape((2,)), self.orig_f(t, y).reshape((2,))) * self.orig_f(t, y))   \n",
    "\n",
    "    def forward(self, t, y):\n",
    "        return (self.orig_f(t, y) * (1/torch.norm(self.orig_f(t, y))**self.beta) \n",
    "                * torch.norm(self.sec_order_deriv(t, y))**self.gamma)*self.const\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2d7bd410",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "generating the   1-the spiral now, remaining time = 263.48 s\n",
      "generating the  51-the spiral now, remaining time = 282.36 s\n",
      "generating the 101-the spiral now, remaining time = 251.14 s\n",
      "generating the 151-the spiral now, remaining time = 216.07 s\n",
      "generating the 201-the spiral now, remaining time = 191.26 s\n",
      "generating the 251-the spiral now, remaining time = 157.35 s\n",
      "generating the 301-the spiral now, remaining time = 122.57 s\n",
      "generating the 351-the spiral now, remaining time = 89.91 s\n",
      "generating the 401-the spiral now, remaining time = 59.62 s\n",
      "generating the 451-the spiral now, remaining time = 29.64 s\n",
      "generating spirals takes 301.89 s\n"
     ]
    }
   ],
   "source": [
    "# generate spirals\n",
    "num_traj   = 500\n",
    "num_points_dense = 400\n",
    "num_points_dense_discard = 0\n",
    "\n",
    "alpha = 2\n",
    "beta = 0\n",
    "gamma = 0\n",
    "const = 1\n",
    "\n",
    "\n",
    "orig_ts_dense = torch.linspace(0,40,num_points_dense+1)\n",
    "# orig_ts_dense = torch.linspace(60., 160., num_points_dense+1)\n",
    "orig_samps_dense = []\n",
    "\n",
    "true_A_mean = torch.tensor([[-0.1, -1.5], [1.5, -0.1]])#torch.tensor([[-0.1, -2], [2, -0.1]])\n",
    "true_A_std = torch.tensor([[0.02, 0.04], [0.04, 0.02]])\n",
    "true_y0 = torch.tensor([[0.6, 0.3]])\n",
    "true_A_list = []\n",
    "\n",
    "t1 = time.time()\n",
    "for i in range(1,num_traj+1):\n",
    "    true_A = torch.tensor(np.float32(np.random.normal(true_A_mean, true_A_std)))\n",
    "    true_A_list.append(true_A)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        orig_samps_temp = odeint(Lambda(-true_A, alpha=alpha, beta=beta, gamma=gamma, const=const), true_y0, orig_ts_dense, method='dopri5')[num_points_dense_discard:,0,:].reshape(\n",
    "            (num_points_dense-num_points_dense_discard+1,2))\n",
    "    orig_samps_dense.append(orig_samps_temp)\n",
    "    if i % 50 == 1:\n",
    "        t3 = time.time()\n",
    "        remain_time = (t3-t1) * (num_traj-i) / i\n",
    "        print('generating the %3d-the spiral now, remaining time = %5.2f s' % (i, remain_time))\n",
    "t2 = time.time()\n",
    "print('generating spirals takes %.2f s' % (t2-t1))\n",
    "    \n",
    "orig_samps_dense = torch.stack(orig_samps_dense).reshape((num_traj, num_points_dense-num_points_dense_discard+1, 2))\n",
    "orig_ts_dense = orig_ts_dense[num_points_dense_discard:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ed8c040b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(401,)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_points = 400\n",
    "subsamp_ind = np.linspace(0, num_points_dense-num_points_dense_discard, num_points+1, dtype=int)\n",
    "orig_samps = orig_samps_dense[:,subsamp_ind,:]\n",
    "orig_ts = orig_ts_dense[subsamp_ind]\n",
    "\n",
    "subsamp_ind.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa362535",
   "metadata": {},
   "source": [
    "# Generate training and testing windows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dea1d9dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_fine_grid = 64\n",
    "num_layers = 3\n",
    "num_coarse_grid = num_fine_grid // 2**(num_layers-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dcfa58f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed_everything()\n",
    "rng = np.random.default_rng() \n",
    "\n",
    "obs_dim = 2\n",
    "num_train_windows = 500\n",
    "num_rep = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "79c709b3",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def sample_windows(trajs_orig, ts_orig, window_length=num_fine_grid+1, num_trajs=250, num_window_per_traj=20,\n",
    "                  ind_start=2, ind_end=200):\n",
    "    train_windows0 = []\n",
    "    train_ts0 = []\n",
    "    train_start0 = []\n",
    "\n",
    "    train_ind_list_candi = np.arange(ind_start,ind_end-window_length)\n",
    "    for j in range(num_trajs):\n",
    "        train_ind_list = np.sort(rng.choice(train_ind_list_candi, num_window_per_traj))\n",
    "        train_start0.append(train_ind_list)\n",
    "        for i in range(num_window_per_traj):\n",
    "            start = train_ind_list[i]\n",
    "            train_windows0.append(torch.from_numpy(np.float32(\n",
    "                trajs_orig[j,start:start+window_length,:])))\n",
    "            train_ts0.append(torch.from_numpy(np.float32(ts_orig[start:start+window_length])))        \n",
    "\n",
    "    train_windows0 = torch.stack(train_windows0)\n",
    "    train_ts0 = torch.stack(train_ts0)\n",
    "\n",
    "    return train_windows0, train_ts0, train_start0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "810bbfd6",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# generate windows\n",
    "\n",
    "train_windows0_list, train_ts0_list, train_start0_list = [], [], []\n",
    "test_windows0_list, test_ts0_list, test_start0_list = [], [], []\n",
    "\n",
    "num_traj_train = 250\n",
    "num_window_per_traj = 2\n",
    "\n",
    "for rep in range(num_rep):\n",
    "    train_windows0, train_ts0, train_start0 = sample_windows(\n",
    "        orig_samps, orig_ts, window_length=num_fine_grid+1, num_trajs=num_traj_train, num_window_per_traj=num_window_per_traj,\n",
    "          ind_start=2, ind_end=num_points)\n",
    "#             print(train_windows0.shape)\n",
    "    test_windows0, test_ts0, test_start0 = sample_windows(\n",
    "        orig_samps[num_traj_train:], orig_ts, window_length=num_fine_grid+1, num_trajs=num_traj_train, num_window_per_traj=num_window_per_traj,\n",
    "          ind_start=2, ind_end=num_points)\n",
    "    train_windows0_list.append(train_windows0)\n",
    "    train_ts0_list.append(train_ts0)\n",
    "    train_start0_list.append(train_start0)\n",
    "    test_windows0_list.append(test_windows0)\n",
    "    test_ts0_list.append(test_ts0)\n",
    "    test_start0_list.append(test_start0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6b3d8a3b",
   "metadata": {
    "code_folding": [
     0,
     1
    ]
   },
   "outputs": [],
   "source": [
    "# customize dataset used in training models\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, train_windows, train_ts):\n",
    "        self.train_windows = train_windows\n",
    "        self.train_ts = train_ts\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        window_temp = self.train_windows[index]\n",
    "        window_ts_temp = self.train_ts[index]\n",
    "        \n",
    "        return window_temp, window_ts_temp\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.train_windows.shape[0]\n",
    "    \n",
    "class MyDataset_adap(Dataset):\n",
    "    def __init__(self, train_windows, train_ts, train_data_len):\n",
    "        self.train_windows = train_windows\n",
    "        self.train_ts = train_ts\n",
    "        self.train_data_len = train_data_len\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        window_temp = self.train_windows[index]\n",
    "        window_ts_temp = self.train_ts[index]\n",
    "        window_temp_len = self.train_data_len[index]\n",
    "        \n",
    "        return window_temp, window_ts_temp, window_temp_len\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.train_windows.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c34937f2",
   "metadata": {},
   "source": [
    "# RNN_ODE models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dadefef",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# RNN_ODE, number of grids >= 16\n",
    "\n",
    "\n",
    "num_grids_non_adap_list1 = np.arange(16,68,4, dtype=int)+1#np.arange(24,102,6, dtype=int)+1#\n",
    "\n",
    "odefunc_non_adap_list1, outputfunc_non_adap_list1 = [], []\n",
    "fit_window_train_non_adap_list1, fit_window_test_non_adap_list1 = [], []\n",
    "fit_L2_err_train_non_adap_list1, fit_L2_err_test_non_adap_list1 = [], []\n",
    "fit_L1_err_train_non_adap_list1, fit_L1_err_test_non_adap_list1 = [], []\n",
    "\n",
    "fit_L2_err_train_non_adap_mean_list1, fit_L2_err_test_non_adap_mean_list1 = [], []\n",
    "fit_L2_err_train_non_adap_std_list1, fit_L2_err_test_non_adap_std_list1 = [], []\n",
    "\n",
    "\n",
    "pred_window_train_non_adap_list1, pred_window_test_non_adap_list1 = [], []\n",
    "pred_window_train_non_adap_list1_pre, pred_window_test_non_adap_list1_pre = [], []\n",
    "pred_L2_err_train_non_adap_list1, pred_L2_err_test_non_adap_list1 = [], []\n",
    "pred_L1_err_train_non_adap_list1, pred_L1_err_test_non_adap_list1 = [], []\n",
    "\n",
    "pred_L2_err_train_non_adap_mean_list1, pred_L2_err_test_non_adap_mean_list1 = [], []\n",
    "pred_L2_err_train_non_adap_std_list1, pred_L2_err_test_non_adap_std_list1 = [], []\n",
    "\n",
    "    \n",
    "batch_size = 64#128\n",
    "verbose = 2\n",
    "method = \"naiveEuler\"#\"rk4\"\n",
    "buffer_start_steps = 2\n",
    "n_iter = 40\n",
    "rescale_const = 1\n",
    "\n",
    "t1 = time.time()\n",
    "index = np.arange(num_fine_grid+1)\n",
    "for i, num_grids in enumerate(num_grids_non_adap_list1):\n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "    print('\\nnon-adaptive method, number of grids = %d\\n' % num_grids)\n",
    "    index1 = np.round(np.arange(num_grids-1) * num_fine_grid/(num_grids-1))\n",
    "    index1 = np.concatenate(([index1[0]]*buffer_start_steps, index1, [num_fine_grid])) \n",
    "\n",
    "    print('train models:\\n')\n",
    "\n",
    "    odefunc1_list, outputfunc1_list = [], []\n",
    "    fit_window_train1_list, fit_window_test1_list = [], []\n",
    "    fit_L2_err_train1_list, fit_L2_err_test1_list = [], []\n",
    "    fit_L1_err_train1_list, fit_L1_err_test1_list = [], []\n",
    "\n",
    "    pred_window_train1_list, pred_window_test1_list = [], []\n",
    "    pred_window_train1_list_pre, pred_window_test1_list_pre = [], []\n",
    "    pred_L2_err_train1_list, pred_L2_err_test1_list = [], []\n",
    "    pred_L1_err_train1_list, pred_L1_err_test1_list = [], []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "        \n",
    "        train_ts1 = torch.tensor(np.linspace(train_ts0[:,0], train_ts0[:,-1], num_grids).T)\n",
    "        train_ts1 = torch.concat((train_ts1[:,0].reshape(train_ts0.shape[0],1).repeat(1,buffer_start_steps), \n",
    "                                train_ts1),axis=1)\n",
    "        train_windows1 = interpolate_window(train_windows0, train_ts0, train_ts1, 'cubic')\n",
    "        \n",
    "        if buffer_start_steps > 0:\n",
    "            deltat = train_ts1[0][-1] - train_ts1[0][-2]\n",
    "            train_ts1[:,:buffer_start_steps] = (train_ts1[:,:buffer_start_steps] + \n",
    "                                                torch.arange(-deltat*buffer_start_steps, 0, deltat))\n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "\n",
    "        train_loader1 = torch.utils.data.DataLoader(\n",
    "            MyDataset(train_windows=train_windows1[train_index,:,:], train_ts=train_ts1[train_index,:]), \n",
    "            shuffle=True, batch_size=batch_size,) # rescale the data\n",
    "\n",
    "        input_steps1 = train_windows1.shape[1]  \n",
    "\n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "        flag, odefunc1, outputfunc1 = train_non_adap_models(train_loader1, input_steps1, valid_windows0, valid_ts0, num_grids,  \n",
    "                                                            verbose, n_iter=n_iter, method=method, thres1=6.5e2, thres2=6.5e2, weight=(1,0),\n",
    "                                                            obs_dim=obs_dim, n_hidden=n_hidden, n_latent=n_latent, num_train_windows=500, \n",
    "                                                            time_scale=time_scale, buffer_start_steps=buffer_start_steps)\n",
    "    \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "        \n",
    "        print('\\nfitting error of training data:')\n",
    "        fit_window_train1, fit_L2_err_train1, fit_L1_err_train1 = fit_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, train_windows0, train_ts0, num_grids, method=method, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "        print('\\nfitting error of testing data:')\n",
    "        fit_window_test1, fit_L2_err_test1, fit_L1_err_test1 = fit_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, test_windows0, test_ts0, num_grids, method=method, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "\n",
    "        print('\\npredicting error of training data:')\n",
    "        pred_window_train1_pre, pred_window_train1, pred_L2_err_train1, pred_L1_err_train1 = pred_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "        print('\\npredicting error of testing data:')\n",
    "        pred_window_test1_pre, pred_window_test1, pred_L2_err_test1, pred_L1_err_test1 = pred_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "\n",
    "        odefunc1_list.append(copy.deepcopy(odefunc1))\n",
    "        outputfunc1_list.append(copy.deepcopy(outputfunc1))\n",
    "        fit_L2_err_train1_list.append(fit_L2_err_train1)\n",
    "        fit_L2_err_test1_list.append(fit_L2_err_test1)\n",
    "        fit_L1_err_train1_list.append(fit_L1_err_train1)\n",
    "        fit_L1_err_test1_list.append(fit_L1_err_test1)\n",
    "        fit_window_train1_list.append(fit_window_train1)\n",
    "        fit_window_test1_list.append(fit_window_test1)\n",
    "\n",
    "        pred_L2_err_train1_list.append(pred_L2_err_train1)\n",
    "        pred_L2_err_test1_list.append(pred_L2_err_test1)\n",
    "        pred_L1_err_train1_list.append(pred_L1_err_train1)\n",
    "        pred_L1_err_test1_list.append(pred_L1_err_test1)\n",
    "        pred_window_train1_list.append(pred_window_train1)\n",
    "        pred_window_test1_list.append(pred_window_test1)\n",
    "        pred_window_train1_list_pre.append(pred_window_train1_pre)\n",
    "        pred_window_test1_list_pre.append(pred_window_test1_pre)\n",
    "\n",
    "    odefunc_non_adap_list1.append(odefunc1_list)\n",
    "    outputfunc_non_adap_list1.append(outputfunc1_list)\n",
    "    fit_window_train_non_adap_list1.append(fit_window_train1_list)\n",
    "    fit_window_test_non_adap_list1.append(fit_window_test1_list)\n",
    "    fit_L2_err_train_non_adap_list1.append(fit_L2_err_train1_list)\n",
    "    fit_L2_err_test_non_adap_list1.append(fit_L2_err_test1_list)\n",
    "    fit_L1_err_train_non_adap_list1.append(fit_L1_err_train1_list)\n",
    "    fit_L1_err_test_non_adap_list1.append(fit_L1_err_test1_list)    \n",
    "    \n",
    "    fit_L2_err_train_ave = np.mean(torch.stack(fit_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave = np.mean(torch.stack(fit_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_train_non_adap_mean_list1.append(np.mean(fit_L2_err_train_ave))\n",
    "    fit_L2_err_train_non_adap_std_list1.append(np.std(fit_L2_err_train_ave))\n",
    "    fit_L2_err_test_non_adap_mean_list1.append(np.mean(fit_L2_err_test_ave))\n",
    "    fit_L2_err_test_non_adap_std_list1.append(np.std(fit_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave),\n",
    "                                                                         np.std(fit_L2_err_train_ave)))\n",
    "    print('\\nL2 fitting error of testing data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_test_ave),\n",
    "                                                                         np.std(fit_L2_err_test_ave)))\n",
    "    \n",
    "    pred_window_train_non_adap_list1.append(pred_window_train1_list)\n",
    "    pred_window_test_non_adap_list1.append(pred_window_test1_list)\n",
    "    pred_window_train_non_adap_list1_pre.append(pred_window_train1_list_pre)\n",
    "    pred_window_test_non_adap_list1_pre.append(pred_window_test1_list_pre)\n",
    "    pred_L2_err_train_non_adap_list1.append(pred_L2_err_train1_list)\n",
    "    pred_L2_err_test_non_adap_list1.append(pred_L2_err_test1_list)\n",
    "    pred_L1_err_train_non_adap_list1.append(pred_L1_err_train1_list)\n",
    "    pred_L1_err_test_non_adap_list1.append(pred_L1_err_test1_list)    \n",
    "    \n",
    "    pred_L2_err_train_ave = np.mean(torch.stack(pred_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_test_ave = np.mean(torch.stack(pred_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_train_non_adap_mean_list1.append(np.mean(pred_L2_err_train_ave))\n",
    "    pred_L2_err_train_non_adap_std_list1.append(np.std(pred_L2_err_train_ave))\n",
    "    pred_L2_err_test_non_adap_mean_list1.append(np.mean(pred_L2_err_test_ave))\n",
    "    pred_L2_err_test_non_adap_std_list1.append(np.std(pred_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 predicting error of training data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_train_ave),\n",
    "                                                                         np.std(pred_L2_err_train_ave)))\n",
    "    print('\\nL2 predicting error of testing data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave),\n",
    "                                                                         np.std(pred_L2_err_test_ave)))\n",
    "    \n",
    "t2 = time.time()\n",
    "print('training model takes %.2f s' % (t2-t1) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ed4265",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# non-adaptive method, interpolate to get the regular time steps (validation data), smaller number of grids\n",
    "\n",
    "\n",
    "num_grids_non_adap_list2 = np.array([6,8,12])+1#np.arange(24,102,6, dtype=int)+1#\n",
    "\n",
    "odefunc_non_adap_list2, outputfunc_non_adap_list2 = [], []\n",
    "fit_window_train_non_adap_list2, fit_window_test_non_adap_list2 = [], []\n",
    "fit_L2_err_train_non_adap_list2, fit_L2_err_test_non_adap_list2 = [], []\n",
    "fit_L1_err_train_non_adap_list2, fit_L1_err_test_non_adap_list2 = [], []\n",
    "\n",
    "fit_L2_err_train_non_adap_mean_list2, fit_L2_err_test_non_adap_mean_list2 = [], []\n",
    "fit_L2_err_train_non_adap_std_list2, fit_L2_err_test_non_adap_std_list2 = [], []\n",
    "\n",
    "\n",
    "pred_window_train_non_adap_list2, pred_window_test_non_adap_list2 = [], []\n",
    "pred_window_train_non_adap_list2_pre, pred_window_test_non_adap_list2_pre = [], []\n",
    "pred_L2_err_train_non_adap_list2, pred_L2_err_test_non_adap_list2 = [], []\n",
    "pred_L1_err_train_non_adap_list2, pred_L1_err_test_non_adap_list2 = [], []\n",
    "\n",
    "pred_L2_err_train_non_adap_mean_list2, pred_L2_err_test_non_adap_mean_list2 = [], []\n",
    "pred_L2_err_train_non_adap_std_list2, pred_L2_err_test_non_adap_std_list2 = [], []\n",
    "\n",
    "    \n",
    "\n",
    "batch_size = 64#128\n",
    "verbose = 2\n",
    "method = \"naiveEuler\"#\"rk4\"\n",
    "buffer_start_steps = 2\n",
    "\n",
    "t1 = time.time()\n",
    "index = np.arange(num_fine_grid+1)\n",
    "for i, num_grids in enumerate(num_grids_non_adap_list2):\n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "    print('\\nnon-adaptive method, number of grids = %d\\n' % num_grids)\n",
    "    index1 = np.round(np.arange(num_grids-1) * num_fine_grid/(num_grids-1))\n",
    "    index1 = np.concatenate(([index1[0]]*buffer_start_steps, index1, [num_fine_grid])) \n",
    "\n",
    "    print('train models:\\n')\n",
    "\n",
    "    odefunc1_list, outputfunc1_list = [], []\n",
    "    fit_window_train1_list, fit_window_test1_list = [], []\n",
    "    fit_L2_err_train1_list, fit_L2_err_test1_list = [], []\n",
    "    fit_L1_err_train1_list, fit_L1_err_test1_list = [], []\n",
    "\n",
    "    pred_window_train1_list, pred_window_test1_list = [], []\n",
    "    pred_window_train1_list_pre, pred_window_test1_list_pre = [], []\n",
    "    pred_L2_err_train1_list, pred_L2_err_test1_list = [], []\n",
    "    pred_L1_err_train1_list, pred_L1_err_test1_list = [], []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "        \n",
    "        train_ts1 = torch.tensor(np.linspace(train_ts0[:,0], train_ts0[:,-1], num_grids).T)\n",
    "        train_ts1 = torch.concat((train_ts1[:,0].reshape(train_ts0.shape[0],1).repeat(1,buffer_start_steps), \n",
    "                                train_ts1),axis=1)\n",
    "        train_windows1 = interpolate_window(train_windows0, train_ts0, train_ts1, 'cubic')\n",
    "        \n",
    "        if buffer_start_steps > 0:\n",
    "            deltat = train_ts1[0][-1] - train_ts1[0][-2]\n",
    "            train_ts1[:,:buffer_start_steps] = (train_ts1[:,:buffer_start_steps] + \n",
    "                                                torch.arange(-deltat*buffer_start_steps, 0, deltat))\n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "\n",
    "        train_loader1 = torch.utils.data.DataLoader(\n",
    "            MyDataset(train_windows=train_windows1[train_index,:,:], train_ts=train_ts1[train_index,:]), \n",
    "            shuffle=True, batch_size=batch_size,) # rescale the data\n",
    "\n",
    "        input_steps1 = train_windows1.shape[1]  \n",
    "\n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "        flag, odefunc1, outputfunc1 = train_non_adap_models(train_loader1, input_steps1, valid_windows0, valid_ts0, num_grids,  \n",
    "                                                            verbose, n_iter=n_iter, method=method, thres1=6.5e2, thres2=6.5e2, weight=(1,0),\n",
    "                                                            obs_dim=obs_dim, n_hidden=n_hidden, n_latent=n_latent, num_train_windows=500, \n",
    "                                                            time_scale=time_scale, buffer_start_steps=buffer_start_steps)\n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "        \n",
    "        print('\\nfitting error of training data:')\n",
    "        fit_window_train1, fit_L2_err_train1, fit_L1_err_train1 = fit_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, train_windows0, train_ts0, num_grids, method=method, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "        print('\\nfitting error of testing data:')\n",
    "        fit_window_test1, fit_L2_err_test1, fit_L1_err_test1 = fit_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, test_windows0, test_ts0, num_grids, method=method, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "\n",
    "        print('\\npredicting error of training data:')\n",
    "        pred_window_train1_pre, pred_window_train1, pred_L2_err_train1, pred_L1_err_train1 = pred_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "        print('\\npredicting error of testing data:')\n",
    "        pred_window_test1_pre, pred_window_test1, pred_L2_err_test1, pred_L1_err_test1 = pred_non_adap_models_grids(\n",
    "            odefunc1, outputfunc1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, rescale_const=rescale_const, time_scale=time_scale)\n",
    "\n",
    "        odefunc1_list.append(copy.deepcopy(odefunc1))\n",
    "        outputfunc1_list.append(copy.deepcopy(outputfunc1))\n",
    "        fit_L2_err_train1_list.append(fit_L2_err_train1)\n",
    "        fit_L2_err_test1_list.append(fit_L2_err_test1)\n",
    "        fit_L1_err_train1_list.append(fit_L1_err_train1)\n",
    "        fit_L1_err_test1_list.append(fit_L1_err_test1)\n",
    "        fit_window_train1_list.append(fit_window_train1)\n",
    "        fit_window_test1_list.append(fit_window_test1)\n",
    "\n",
    "        pred_L2_err_train1_list.append(pred_L2_err_train1)\n",
    "        pred_L2_err_test1_list.append(pred_L2_err_test1)\n",
    "        pred_L1_err_train1_list.append(pred_L1_err_train1)\n",
    "        pred_L1_err_test1_list.append(pred_L1_err_test1)\n",
    "        pred_window_train1_list.append(pred_window_train1)\n",
    "        pred_window_test1_list.append(pred_window_test1)\n",
    "        pred_window_train1_list_pre.append(pred_window_train1_pre)\n",
    "        pred_window_test1_list_pre.append(pred_window_test1_pre)\n",
    "\n",
    "    odefunc_non_adap_list2.append(odefunc1_list)\n",
    "    outputfunc_non_adap_list2.append(outputfunc1_list)\n",
    "    fit_window_train_non_adap_list2.append(fit_window_train1_list)\n",
    "    fit_window_test_non_adap_list2.append(fit_window_test1_list)\n",
    "    fit_L2_err_train_non_adap_list2.append(fit_L2_err_train1_list)\n",
    "    fit_L2_err_test_non_adap_list2.append(fit_L2_err_test1_list)\n",
    "    fit_L1_err_train_non_adap_list2.append(fit_L1_err_train1_list)\n",
    "    fit_L1_err_test_non_adap_list2.append(fit_L1_err_test1_list)    \n",
    "    \n",
    "    fit_L2_err_train_ave = np.mean(torch.stack(fit_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave = np.mean(torch.stack(fit_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_train_non_adap_mean_list2.append(np.mean(fit_L2_err_train_ave))\n",
    "    fit_L2_err_train_non_adap_std_list2.append(np.std(fit_L2_err_train_ave))\n",
    "    fit_L2_err_test_non_adap_mean_list2.append(np.mean(fit_L2_err_test_ave))\n",
    "    fit_L2_err_test_non_adap_std_list2.append(np.std(fit_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave),\n",
    "                                                                         np.std(fit_L2_err_train_ave)))\n",
    "    print('\\nL2 fitting error of testing data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_test_ave),\n",
    "                                                                         np.std(fit_L2_err_test_ave)))\n",
    "    \n",
    "    pred_window_train_non_adap_list2.append(pred_window_train1_list)\n",
    "    pred_window_test_non_adap_list2.append(pred_window_test1_list)\n",
    "    pred_window_train_non_adap_list2_pre.append(pred_window_train1_list_pre)\n",
    "    pred_window_test_non_adap_list2_pre.append(pred_window_test1_list_pre)\n",
    "    pred_L2_err_train_non_adap_list2.append(pred_L2_err_train1_list)\n",
    "    pred_L2_err_test_non_adap_list2.append(pred_L2_err_test1_list)\n",
    "    pred_L1_err_train_non_adap_list2.append(pred_L1_err_train1_list)\n",
    "    pred_L1_err_test_non_adap_list2.append(pred_L1_err_test1_list)    \n",
    "    \n",
    "    pred_L2_err_train_ave = np.mean(torch.stack(pred_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_test_ave = np.mean(torch.stack(pred_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_train_non_adap_mean_list2.append(np.mean(pred_L2_err_train_ave))\n",
    "    pred_L2_err_train_non_adap_std_list2.append(np.std(pred_L2_err_train_ave))\n",
    "    pred_L2_err_test_non_adap_mean_list2.append(np.mean(pred_L2_err_test_ave))\n",
    "    pred_L2_err_test_non_adap_std_list2.append(np.std(pred_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 predicting error of training data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_train_ave),\n",
    "                                                                         np.std(pred_L2_err_train_ave)))\n",
    "    print('\\nL2 predicting error of testing data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave),\n",
    "                                                                         np.std(pred_L2_err_test_ave)))\n",
    "    \n",
    "t2 = time.time()\n",
    "print('training model takes %.2f s' % (t2-t1) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "448acf10",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# compute errors\n",
    "\n",
    "pred_L2_err_test_non_adap_list_mean_mat1 = np.zeros((num_rep, len(num_grids_non_adap_list1)))\n",
    "\n",
    "for i in range(len(num_grids_non_adap_list1)):\n",
    "    pred_L2_err_test_non_adap_list_mean_mat1[:,i] = torch.stack(pred_L2_err_test_non_adap_list1[i]).mean(axis=(1,)).numpy()\n",
    "\n",
    "    \n",
    "pred_L2_err_test_non_adap_list_mean_mat2 = np.zeros((num_rep, len(num_grids_non_adap_list2)))\n",
    "\n",
    "for i in range(len(num_grids_non_adap_list2)):\n",
    "    pred_L2_err_test_non_adap_list_mean_mat2[:,i] = torch.stack(pred_L2_err_test_non_adap_list2[i]).mean(axis=(1,)).numpy()\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96271747",
   "metadata": {},
   "source": [
    "# RNN_ODE_Adap models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82cd755d",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# RNN_ODE_Adap, number of grids >= 16\n",
    "\n",
    "thres_list1 = ([0, 0.565, 0.615, 0.66, 0.73, 0.81, 0.91, 1.09] + \n",
    "              [0.84, 0.91, 1.02, 1.18, 3])\n",
    "\n",
    "rescale_const = 1\n",
    "thres_list1.reverse()\n",
    "\n",
    "\n",
    "verbose = 2\n",
    "n_iter = 50\n",
    "\n",
    "num_grids_adap_list1 = []\n",
    "\n",
    "odefunc_adap_list1, outputfunc_adap_list1 = [], []\n",
    "\n",
    "# fitting errors\n",
    "fit_window_train_adap_list1_cubic, fit_window_test_adap_list1_cubic = [], []\n",
    "fit_L2_err_train_adap_list1_cubic, fit_L2_err_test_adap_list1_cubic = [], []\n",
    "fit_L1_err_train_adap_list1_cubic, fit_L1_err_test_adap_list1_cubic = [], []\n",
    "\n",
    "fit_L2_err_train_adap_mean_list1_cubic, fit_L2_err_test_adap_mean_list1_cubic = [], []\n",
    "fit_L2_err_train_adap_std_list1_cubic, fit_L2_err_test_adap_std_list1_cubic = [], []\n",
    "\n",
    "\n",
    "# predicting errors\n",
    "pred_window_test_adap_list1_cubic = []\n",
    "pred_window_test_adap_list1_cubic_pre = []\n",
    "pred_L2_err_test_adap_list1_cubic = []\n",
    "pred_L1_err_test_adap_list1_cubic = []\n",
    "\n",
    "pred_L2_err_test_adap_mean_list1_cubic = []\n",
    "pred_L2_err_test_adap_std_list1_cubic = []\n",
    "\n",
    "\n",
    "for i, thres in enumerate(thres_list1):\n",
    "        \n",
    "#     if i <= 1:\n",
    "#         num_fine_adap = 64\n",
    "#         num_layers_adap = 4        \n",
    "#     elif i <= 4:\n",
    "#         num_fine_adap = 64\n",
    "#         num_layers_adap = 3\n",
    "#     else:\n",
    "#         num_fine_adap = 64\n",
    "#         num_layers_adap = 2\n",
    "    if i <= 4:\n",
    "        num_fine_adap = 64\n",
    "        num_layers_adap = 3       \n",
    "    else:\n",
    "        num_fine_adap = 64\n",
    "        num_layers_adap = 2\n",
    "        \n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "\n",
    "    odefunc1_list1, outputfunc1_list1 = [], []\n",
    "    \n",
    "    fit_window_train1_list1_cubic, fit_window_test1_list1_cubic = [], []\n",
    "    fit_L2_err_train1_list1_cubic, fit_L2_err_test1_list1_cubic = [], []\n",
    "    fit_L1_err_train1_list1_cubic, fit_L1_err_test1_list1_cubic = [], []\n",
    "    fit_window_train1_list1_linear, fit_window_test1_list1_linear = [], []\n",
    "    fit_L2_err_train1_list1_linear, fit_L2_err_test1_list1_linear = [], []\n",
    "    fit_L1_err_train1_list1_linear, fit_L1_err_test1_list1_linear = [], []\n",
    "\n",
    "    pred_window_test1_list1_cubic = []\n",
    "    pred_window_test1_list1_cubic_pre = []\n",
    "    pred_L2_err_test1_list1_cubic = []\n",
    "    pred_L1_err_test1_list1_cubic = []\n",
    "    pred_window_test1_list1_linear = []\n",
    "    pred_window_test1_list1_linear_pre = []\n",
    "    pred_L2_err_test1_list1_linear = []\n",
    "    pred_L1_err_test1_list1_linear = []\n",
    "    \n",
    "    num_grids_adap_list1_temp = []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "\n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "\n",
    "        train_windows0_adap_buffer, train_ts0_adap_buffer, train_len0_adap_buffer, train_windows_adap_len = create_adap_data_buffer(\n",
    "            train_windows0, train_ts0, thres, num_fine_adap, num_layers_adap)\n",
    "        test_windows0_adap_buffer, test_ts0_adap_buffer, test_len0_adap_buffer, _ = create_adap_data_buffer(\n",
    "            test_windows0, test_ts0, thres, num_fine_adap, num_layers_adap)\n",
    "        # train_windows, 1st half\n",
    "        train_windows0_adap_buffer_trunc, train_ts0_adap_buffer_trunc, _, _, train_time_index = create_adap_data_buffer(\n",
    "            train_windows0[:,:(train_ts0.shape[1]+1)//2,:], train_ts0[:,:(train_ts0.shape[1]+1)//2], thres, num_fine_adap//2, num_layers_adap, return_index=1)\n",
    "        # train_windows, 2nd half\n",
    "        _, train_ts0_adap_buffer_trunc2, train_len0_adap_buffer_trunc2, _ = create_adap_data_buffer(\n",
    "            train_windows0[:,-(train_ts0.shape[1]+1)//2:,:], train_ts0[:,-(train_ts0.shape[1]+1)//2:], thres, num_fine_adap//2, num_layers_adap, buffer_start_steps=0)\n",
    "        # test_windows, 1st half\n",
    "        test_windows0_adap_buffer_trunc, test_ts0_adap_buffer_trunc, _, _, test_time_index = create_adap_data_buffer(\n",
    "            test_windows0[:,:(test_ts0.shape[1]+1)//2,:], test_ts0[:,:(test_ts0.shape[1]+1)//2], thres, num_fine_adap//2, num_layers_adap, return_index=1)\n",
    "        # test_windows, 2nd half\n",
    "        _, test_ts0_adap_buffer_trunc2, test_len0_adap_buffer_trunc2, _ = create_adap_data_buffer(\n",
    "            test_windows0[:,-(test_ts0.shape[1]+1)//2:,:], test_ts0[:,-(test_ts0.shape[1]+1)//2:], thres, num_fine_adap//2, num_layers_adap, buffer_start_steps=0)\n",
    "\n",
    "        num_grids_adap_list1_temp.append(train_windows_adap_len)\n",
    "        \n",
    "#         print(train_windows0[:,:(train_ts0.shape[1]+1)//2,:].shape)\n",
    "\n",
    "        print('\\nadaptive method, thres = %.2f, number of grids = %.1f (%.2e)\\n' % (thres, train_windows_adap_len, train_windows_adap_len))\n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        valid_windows0_adap_buffer = train_windows0_adap_buffer[valid_index,:,:]\n",
    "        valid_ts0_adap_buffer = train_ts0_adap_buffer[valid_index,:]\n",
    "        valid_windows0_adap_buffer_trunc = train_windows0_adap_buffer_trunc[valid_index,:,:]\n",
    "        valid_ts0_adap_buffer_trunc = train_ts0_adap_buffer_trunc[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "        train_loader_adap1_buffer = torch.utils.data.DataLoader(\n",
    "            MyDataset_adap(train_windows=train_windows0_adap_buffer[train_index], train_ts=train_ts0_adap_buffer[train_index], \n",
    "                            train_data_len=np.array(train_len0_adap_buffer)[train_index]), \n",
    "            shuffle=True, batch_size=batch_size,) \n",
    "        input_steps_adap1_buffer = max(train_len0_adap_buffer)   \n",
    "        \n",
    "        print('train models:\\n')\n",
    "        \n",
    "        if i <= 1:\n",
    "            valid_ts_adap_trunc2, valid_len_adap_trunc2 = valid_ts0[:, (valid_ts0.shape[1]-1)//2::1], [(valid_ts0.shape[1]-1)//2+1]*test_ts0.shape[0]\n",
    "        elif i <= 4:\n",
    "            valid_ts_adap_trunc2, valid_len_adap_trunc2 = valid_ts0[:, (valid_ts0.shape[1]-1)//2::1], [(valid_ts0.shape[1]-1)//2+1]*test_ts0.shape[0]\n",
    "        else:\n",
    "            valid_ts_adap_trunc2, valid_len_adap_trunc2 = valid_ts0[:, (valid_ts0.shape[1]-1)//2::1], [(valid_ts0.shape[1]-1)//2+1]*test_ts0.shape[0]\n",
    "\n",
    "        flag, odefunc_adap1_buffer, outputfunc_adap1_buffer = train_adap_models(\n",
    "            train_loader_adap1_buffer, input_steps_adap1_buffer, verbose, \n",
    "            valid_window=valid_windows0, valid_ts=valid_ts0, valid_window_adap=valid_windows0_adap_buffer, \n",
    "            valid_ts_adap=valid_ts0_adap_buffer, valid_window_adap_trunc=valid_windows0_adap_buffer_trunc, \n",
    "            valid_ts_adap_trunc=valid_ts0_adap_buffer_trunc, buffer_start_steps=buffer_start_steps, \n",
    "            weight=(1,0), n_iter=n_iter, thres1=1.5e3, valid_len=np.array(train_len0_adap_buffer)[valid_index], \n",
    "            valid_ts_adap_trunc2=valid_ts_adap_trunc2, valid_len_adap_trunc2=valid_len_adap_trunc2,\n",
    "            pred_len=(train_ts0.shape[1]-1)//2, rescale_const=rescale_const,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )    \n",
    "        \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "            \n",
    "        odefunc1_list1.append(copy.deepcopy(odefunc_adap1_buffer))\n",
    "        outputfunc1_list1.append(copy.deepcopy(outputfunc_adap1_buffer))\n",
    "        \n",
    "        if i <= 1:\n",
    "            test_ts_adap_trunc2, test_len_adap_trunc2 = test_ts0[:, (test_ts0.shape[1]-1)//2::1], [(test_ts0.shape[1]-1)//2+1]*test_ts0.shape[0]\n",
    "        elif i <= 4:\n",
    "            test_ts_adap_trunc2, test_len_adap_trunc2 = test_ts0[:, (test_ts0.shape[1]-1)//2::1], [(test_ts0.shape[1]-1)//2+1]*test_ts0.shape[0]\n",
    "        else:\n",
    "            test_ts_adap_trunc2, test_len_adap_trunc2 = test_ts0[:, (test_ts0.shape[1]-1)//2::1], [(test_ts0.shape[1]-1)//2+1]*test_ts0.shape[0]\n",
    "        print(test_ts_adap_trunc2.shape)\n",
    "        interp_kind = 'cubic'\n",
    "        print('\\nfitting error of training data (cubic):')\n",
    "        fit_window_adap_train1, fit_L2_err_adap_train1, fit_L1_err_adap_train1 = fit_adap_models(\n",
    "            odefunc_adap1_buffer, outputfunc_adap1_buffer, train_windows0, train_ts0, \n",
    "            train_windows0_adap_buffer, train_ts0_adap_buffer, train_len0_adap_buffer, buffer_start_steps,\n",
    "            interp_kind=interp_kind, rescale_const=rescale_const, n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )   \n",
    "        print('fitting error of testing data (cubic):')\n",
    "        fit_window_adap_test1, fit_L2_err_adap_test1, fit_L1_err_adap_test1 = fit_adap_models(\n",
    "            odefunc_adap1_buffer, outputfunc_adap1_buffer, test_windows0, test_ts0, \n",
    "            test_windows0_adap_buffer, test_ts0_adap_buffer, test_len0_adap_buffer, buffer_start_steps,\n",
    "            interp_kind=interp_kind, rescale_const=rescale_const, n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )    \n",
    "        fit_L2_err_train1_list1_cubic.append(fit_L2_err_adap_train1)\n",
    "        fit_L2_err_test1_list1_cubic.append(fit_L2_err_adap_test1)\n",
    "        fit_L1_err_train1_list1_cubic.append(fit_L1_err_adap_train1)\n",
    "        fit_L1_err_test1_list1_cubic.append(fit_L1_err_adap_test1)\n",
    "        fit_window_train1_list1_cubic.append(fit_window_adap_train1)\n",
    "        fit_window_test1_list1_cubic.append(fit_window_adap_test1)\n",
    "\n",
    "        print('predicting error of testing data (cubic):')\n",
    "        pred_window_adap_test1_pre, pred_window_adap_test1, pred_L2_err_adap_test1, pred_L1_err_adap_test1 = pred_adap_models(\n",
    "            odefunc_adap1_buffer, outputfunc_adap1_buffer, test_windows0, test_ts0, \n",
    "            test_windows0_adap_buffer_trunc, test_ts0_adap_buffer_trunc, rescale_const=rescale_const,\n",
    "            ts_adap_buffer_trunc2=test_ts_adap_trunc2, len_adap_buffer_trunc2=test_len_adap_trunc2, \n",
    "            pred_len=(train_ts0.shape[1]-1)//2, buffer_start_steps=buffer_start_steps,interp_kind=interp_kind,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )     \n",
    "\n",
    "        pred_L2_err_test1_list1_cubic.append(pred_L2_err_adap_test1)\n",
    "        pred_L1_err_test1_list1_cubic.append(pred_L1_err_adap_test1)\n",
    "        pred_window_test1_list1_cubic.append(pred_window_adap_test1)\n",
    "        pred_window_test1_list1_cubic_pre.append(pred_window_adap_test1_pre)\n",
    "        \n",
    "        \n",
    "        \n",
    "    odefunc_adap_list1.append(odefunc1_list1)\n",
    "    outputfunc_adap_list1.append(outputfunc1_list1)\n",
    "    num_grids_adap_list1.append(np.mean(num_grids_adap_list1_temp))\n",
    "    \n",
    "    fit_window_train_adap_list1_cubic.append(fit_window_train1_list1_cubic)\n",
    "    fit_window_test_adap_list1_cubic.append(fit_window_test1_list1_cubic)\n",
    "    fit_L2_err_train_adap_list1_cubic.append(fit_L2_err_train1_list1_cubic)\n",
    "    fit_L2_err_test_adap_list1_cubic.append(fit_L2_err_test1_list1_cubic)\n",
    "    fit_L1_err_train_adap_list1_cubic.append(fit_L1_err_train1_list1_cubic)\n",
    "    fit_L1_err_test_adap_list1_cubic.append(fit_L1_err_test1_list1_cubic)    \n",
    "    \n",
    "    fit_L2_err_train_ave_cubic = np.mean(torch.stack(fit_L2_err_train1_list1_cubic).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave_cubic = np.mean(torch.stack(fit_L2_err_test1_list1_cubic).numpy(),axis=(1,))\n",
    "    \n",
    "    fit_L2_err_train_adap_mean_list1_cubic.append(np.mean(fit_L2_err_train_ave_cubic))\n",
    "    fit_L2_err_train_adap_std_list1_cubic.append(np.std(fit_L2_err_train_ave_cubic))\n",
    "    fit_L2_err_test_adap_mean_list1_cubic.append(np.mean(fit_L2_err_test_ave_cubic))\n",
    "    fit_L2_err_test_adap_std_list1_cubic.append(np.std(fit_L2_err_test_ave_cubic))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data (cubic): mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave_cubic),\n",
    "                                                                         np.std(fit_L2_err_train_ave_cubic)))\n",
    "    pred_window_test_adap_list1_cubic.append(pred_window_test1_list1_cubic)\n",
    "    pred_window_test_adap_list1_cubic_pre.append(pred_window_test1_list1_cubic_pre)\n",
    "    pred_L2_err_test_adap_list1_cubic.append(pred_L2_err_test1_list1_cubic)\n",
    "    pred_L1_err_test_adap_list1_cubic.append(pred_L1_err_test1_list1_cubic)    \n",
    "    \n",
    "    pred_L2_err_test_ave_cubic = np.mean(torch.stack(pred_L2_err_test1_list1_cubic).numpy(),axis=(1,))\n",
    "    \n",
    "    pred_L2_err_test_adap_mean_list1_cubic.append(np.mean(pred_L2_err_test_ave_cubic))\n",
    "    pred_L2_err_test_adap_std_list1_cubic.append(np.std(pred_L2_err_test_ave_cubic))\n",
    "    \n",
    "    print('\\nL2 predicting error of testing data (cubic): mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave_cubic),\n",
    "                                                                         np.std(pred_L2_err_test_ave_cubic)))\n",
    "     \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2daa41",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# RNN_ODE_Adap, number of grids < 16\n",
    "\n",
    "thres_list2 = [\n",
    "    0.9, # 8, 2\n",
    "    0, # 8, 2\n",
    "    0.7, # 16, 2\n",
    "              ]\n",
    "\n",
    "rescale_const = 1\n",
    "\n",
    "verbose = 2\n",
    "\n",
    "num_grids_adap_list2 = []\n",
    "\n",
    "odefunc_adap_list2, outputfunc_adap_list2 = [], []\n",
    "\n",
    "# fitting errors\n",
    "fit_window_train_adap_list2_cubic, fit_window_test_adap_list2_cubic = [], []\n",
    "fit_L2_err_train_adap_list2_cubic, fit_L2_err_test_adap_list2_cubic = [], []\n",
    "fit_L1_err_train_adap_list2_cubic, fit_L1_err_test_adap_list2_cubic = [], []\n",
    "\n",
    "fit_L2_err_train_adap_mean_list2_cubic, fit_L2_err_test_adap_mean_list2_cubic = [], []\n",
    "fit_L2_err_train_adap_std_list2_cubic, fit_L2_err_test_adap_std_list2_cubic = [], []\n",
    "\n",
    "\n",
    "# predicting errors\n",
    "pred_window_test_adap_list2_cubic = []\n",
    "pred_window_test_adap_list2_cubic_pre = []\n",
    "pred_L2_err_test_adap_list2_cubic = []\n",
    "pred_L1_err_test_adap_list2_cubic = []\n",
    "\n",
    "pred_L2_err_test_adap_mean_list2_cubic = []\n",
    "pred_L2_err_test_adap_std_list2_cubic = []\n",
    "\n",
    "\n",
    "for i, thres in enumerate(thres_list2):\n",
    "    \n",
    "#     if i <= 1:\n",
    "#         num_fine_adap = 64\n",
    "#         num_layers_adap = 4        \n",
    "#     elif i <= 4:\n",
    "#         num_fine_adap = 64\n",
    "#         num_layers_adap = 3\n",
    "#     else:\n",
    "#         num_fine_adap = 64\n",
    "#         num_layers_adap = 2\n",
    "            \n",
    "    if i <= 1:\n",
    "        num_fine_adap = 8\n",
    "        num_layers_adap = 2\n",
    "    else:\n",
    "        num_fine_adap = 16\n",
    "        num_layers_adap = 2\n",
    "        \n",
    "        \n",
    "        \n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "\n",
    "    odefunc1_list2, outputfunc1_list2 = [], []\n",
    "    \n",
    "    fit_window_train1_list2_cubic, fit_window_test1_list2_cubic = [], []\n",
    "    fit_L2_err_train1_list2_cubic, fit_L2_err_test1_list2_cubic = [], []\n",
    "    fit_L1_err_train1_list2_cubic, fit_L1_err_test1_list2_cubic = [], []\n",
    "    fit_window_train1_list2_linear, fit_window_test1_list2_linear = [], []\n",
    "    fit_L2_err_train1_list2_linear, fit_L2_err_test1_list2_linear = [], []\n",
    "    fit_L1_err_train1_list2_linear, fit_L1_err_test1_list2_linear = [], []\n",
    "\n",
    "    pred_window_test1_list2_cubic = []\n",
    "    pred_window_test1_list2_cubic_pre = []\n",
    "    pred_L2_err_test1_list2_cubic = []\n",
    "    pred_L1_err_test1_list2_cubic = []\n",
    "    pred_window_test1_list2_linear = []\n",
    "    pred_window_test1_list2_linear_pre = []\n",
    "    pred_L2_err_test1_list2_linear = []\n",
    "    pred_L1_err_test1_list2_linear = []\n",
    "    \n",
    "    num_grids_adap_list2_temp = []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "\n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "\n",
    "        train_windows0_adap_buffer, train_ts0_adap_buffer, train_len0_adap_buffer, train_windows_adap_len = create_adap_data_buffer(\n",
    "            train_windows0, train_ts0, thres, num_fine_adap, num_layers_adap)\n",
    "        test_windows0_adap_buffer, test_ts0_adap_buffer, test_len0_adap_buffer, _ = create_adap_data_buffer(\n",
    "            test_windows0, test_ts0, thres, num_fine_adap, num_layers_adap)\n",
    "        # train_windows, 1st half\n",
    "        train_windows0_adap_buffer_trunc, train_ts0_adap_buffer_trunc, _, _, train_time_index = create_adap_data_buffer(\n",
    "            train_windows0[:,:(train_ts0.shape[1]+1)//2,:], train_ts0[:,:(train_ts0.shape[1]+1)//2], thres, num_fine_adap//2, num_layers_adap, return_index=1)\n",
    "        # train_windows, 2nd half\n",
    "        _, train_ts0_adap_buffer_trunc2, train_len0_adap_buffer_trunc2, _ = create_adap_data_buffer(\n",
    "            train_windows0[:,-(train_ts0.shape[1]+1)//2:,:], train_ts0[:,-(train_ts0.shape[1]+1)//2:], thres, num_fine_adap//2, num_layers_adap, buffer_start_steps=0)\n",
    "        # test_windows, 1st half\n",
    "        test_windows0_adap_buffer_trunc, test_ts0_adap_buffer_trunc, _, _, test_time_index = create_adap_data_buffer(\n",
    "            test_windows0[:,:(test_ts0.shape[1]+1)//2,:], test_ts0[:,:(test_ts0.shape[1]+1)//2], thres, num_fine_adap//2, num_layers_adap, return_index=1)\n",
    "        # test_windows, 2nd half\n",
    "        _, test_ts0_adap_buffer_trunc2, test_len0_adap_buffer_trunc2, _ = create_adap_data_buffer(\n",
    "            test_windows0[:,-(test_ts0.shape[1]+1)//2:,:], test_ts0[:,-(test_ts0.shape[1]+1)//2:], thres, num_fine_adap//2, num_layers_adap, buffer_start_steps=0)\n",
    "\n",
    "        num_grids_adap_list2_temp.append(train_windows_adap_len)\n",
    "        \n",
    "#         print(train_windows0[:,:(train_ts0.shape[1]+1)//2,:].shape)\n",
    "\n",
    "        print('\\nadaptive method, thres = %.2f, number of grids = %.1f (%.2e)\\n' % (thres, train_windows_adap_len, train_windows_adap_len))\n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        valid_windows0_adap_buffer = train_windows0_adap_buffer[valid_index,:,:]\n",
    "        valid_ts0_adap_buffer = train_ts0_adap_buffer[valid_index,:]\n",
    "        valid_windows0_adap_buffer_trunc = train_windows0_adap_buffer_trunc[valid_index,:,:]\n",
    "        valid_ts0_adap_buffer_trunc = train_ts0_adap_buffer_trunc[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "        train_loader_adap1_buffer = torch.utils.data.DataLoader(\n",
    "            MyDataset_adap(train_windows=train_windows0_adap_buffer[train_index], train_ts=train_ts0_adap_buffer[train_index], \n",
    "                            train_data_len=np.array(train_len0_adap_buffer)[train_index]), \n",
    "            shuffle=True, batch_size=batch_size,) \n",
    "        input_steps_adap1_buffer = max(train_len0_adap_buffer)   \n",
    "        \n",
    "        print('train models:\\n')\n",
    "        \n",
    "        if i <= 1:\n",
    "            valid_ts_adap_trunc2, valid_len_adap_trunc2 = valid_ts0[:, (valid_ts0.shape[1]-1)//2::8], [(valid_ts0.shape[1]-1)//16+1]*test_ts0.shape[0]\n",
    "        else:\n",
    "            valid_ts_adap_trunc2, valid_len_adap_trunc2 = valid_ts0[:, (valid_ts0.shape[1]-1)//2::4], [(valid_ts0.shape[1]-1)//8+1]*test_ts0.shape[0]\n",
    "\n",
    "        flag, odefunc_adap1_buffer, outputfunc_adap1_buffer = train_adap_models(\n",
    "            train_loader_adap1_buffer, input_steps_adap1_buffer, verbose, \n",
    "            valid_window=valid_windows0, valid_ts=valid_ts0, valid_window_adap=valid_windows0_adap_buffer, \n",
    "            valid_ts_adap=valid_ts0_adap_buffer, valid_window_adap_trunc=valid_windows0_adap_buffer_trunc, \n",
    "            valid_ts_adap_trunc=valid_ts0_adap_buffer_trunc, buffer_start_steps=buffer_start_steps, \n",
    "            weight=[1,0], n_iter=n_iter, thres1=1.5e3, valid_len=np.array(train_len0_adap_buffer)[valid_index], \n",
    "            valid_ts_adap_trunc2=valid_ts_adap_trunc2, valid_len_adap_trunc2=valid_len_adap_trunc2,\n",
    "            pred_len=(train_ts0.shape[1]-1)//2, rescale_const=rescale_const,\n",
    "            n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )    \n",
    "        \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "            \n",
    "        odefunc1_list2.append(copy.deepcopy(odefunc_adap1_buffer))\n",
    "        outputfunc1_list2.append(copy.deepcopy(outputfunc_adap1_buffer))\n",
    "        \n",
    "        if i <= 1:\n",
    "            test_ts_adap_trunc2, test_len_adap_trunc2 = test_ts0[:, (test_ts0.shape[1]-1)//2::8], [(test_ts0.shape[1]-1)//16+1]*test_ts0.shape[0]\n",
    "        else:\n",
    "            test_ts_adap_trunc2, test_len_adap_trunc2 = test_ts0[:, (test_ts0.shape[1]-1)//2::4], [(test_ts0.shape[1]-1)//8+1]*test_ts0.shape[0]\n",
    "        print(test_ts_adap_trunc2.shape)\n",
    "        interp_kind = 'cubic'\n",
    "        print('\\nfitting error of training data (cubic):')\n",
    "        fit_window_adap_train1, fit_L2_err_adap_train1, fit_L1_err_adap_train1 = fit_adap_models(\n",
    "            odefunc_adap1_buffer, outputfunc_adap1_buffer, train_windows0, train_ts0, \n",
    "            train_windows0_adap_buffer, train_ts0_adap_buffer, train_len0_adap_buffer, buffer_start_steps,\n",
    "            interp_kind=interp_kind, rescale_const=rescale_const, n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )   \n",
    "        print('fitting error of testing data (cubic):')\n",
    "        fit_window_adap_test1, fit_L2_err_adap_test1, fit_L1_err_adap_test1 = fit_adap_models(\n",
    "            odefunc_adap1_buffer, outputfunc_adap1_buffer, test_windows0, test_ts0, \n",
    "            test_windows0_adap_buffer, test_ts0_adap_buffer, test_len0_adap_buffer, buffer_start_steps,\n",
    "            interp_kind=interp_kind, rescale_const=rescale_const, n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )    \n",
    "        fit_L2_err_train1_list2_cubic.append(fit_L2_err_adap_train1)\n",
    "        fit_L2_err_test1_list2_cubic.append(fit_L2_err_adap_test1)\n",
    "        fit_L1_err_train1_list2_cubic.append(fit_L1_err_adap_train1)\n",
    "        fit_L1_err_test1_list2_cubic.append(fit_L1_err_adap_test1)\n",
    "        fit_window_train1_list2_cubic.append(fit_window_adap_train1)\n",
    "        fit_window_test1_list2_cubic.append(fit_window_adap_test1)\n",
    "\n",
    "        print('predicting error of testing data (cubic):')\n",
    "        pred_window_adap_test1_pre, pred_window_adap_test1, pred_L2_err_adap_test1, pred_L1_err_adap_test1 = pred_adap_models(\n",
    "            odefunc_adap1_buffer, outputfunc_adap1_buffer, test_windows0, test_ts0, \n",
    "            test_windows0_adap_buffer_trunc, test_ts0_adap_buffer_trunc, rescale_const=rescale_const,\n",
    "            ts_adap_buffer_trunc2=test_ts_adap_trunc2, len_adap_buffer_trunc2=test_len_adap_trunc2, \n",
    "            pred_len=(train_ts0.shape[1]-1)//2, buffer_start_steps=buffer_start_steps,interp_kind=interp_kind, \n",
    "            n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale\n",
    "        )     \n",
    "\n",
    "        pred_L2_err_test1_list2_cubic.append(pred_L2_err_adap_test1)\n",
    "        pred_L1_err_test1_list2_cubic.append(pred_L1_err_adap_test1)\n",
    "        pred_window_test1_list2_cubic.append(pred_window_adap_test1)\n",
    "        pred_window_test1_list2_cubic_pre.append(pred_window_adap_test1_pre)\n",
    "        \n",
    "\n",
    "        \n",
    "    odefunc_adap_list2.append(odefunc1_list2)\n",
    "    outputfunc_adap_list2.append(outputfunc1_list2)\n",
    "    num_grids_adap_list2.append(np.mean(num_grids_adap_list2_temp))\n",
    "    \n",
    "    fit_window_train_adap_list2_cubic.append(fit_window_train1_list2_cubic)\n",
    "    fit_window_test_adap_list2_cubic.append(fit_window_test1_list2_cubic)\n",
    "    fit_L2_err_train_adap_list2_cubic.append(fit_L2_err_train1_list2_cubic)\n",
    "    fit_L2_err_test_adap_list2_cubic.append(fit_L2_err_test1_list2_cubic)\n",
    "    fit_L1_err_train_adap_list2_cubic.append(fit_L1_err_train1_list2_cubic)\n",
    "    fit_L1_err_test_adap_list2_cubic.append(fit_L1_err_test1_list2_cubic)    \n",
    "    \n",
    "    fit_L2_err_train_ave_cubic = np.mean(torch.stack(fit_L2_err_train1_list2_cubic).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave_cubic = np.mean(torch.stack(fit_L2_err_test1_list2_cubic).numpy(),axis=(1,))\n",
    "    \n",
    "    fit_L2_err_train_adap_mean_list2_cubic.append(np.mean(fit_L2_err_train_ave_cubic))\n",
    "    fit_L2_err_train_adap_std_list2_cubic.append(np.std(fit_L2_err_train_ave_cubic))\n",
    "    fit_L2_err_test_adap_mean_list2_cubic.append(np.mean(fit_L2_err_test_ave_cubic))\n",
    "    fit_L2_err_test_adap_std_list2_cubic.append(np.std(fit_L2_err_test_ave_cubic))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data (cubic): mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave_cubic),\n",
    "                                                                         np.std(fit_L2_err_train_ave_cubic)))\n",
    "    pred_window_test_adap_list2_cubic.append(pred_window_test1_list2_cubic)\n",
    "    pred_window_test_adap_list2_cubic_pre.append(pred_window_test1_list2_cubic_pre)\n",
    "    pred_L2_err_test_adap_list2_cubic.append(pred_L2_err_test1_list2_cubic)\n",
    "    pred_L1_err_test_adap_list2_cubic.append(pred_L1_err_test1_list2_cubic)    \n",
    "    \n",
    "    pred_L2_err_test_ave_cubic = np.mean(torch.stack(pred_L2_err_test1_list2_cubic).numpy(),axis=(1,))\n",
    "    \n",
    "    pred_L2_err_test_adap_mean_list2_cubic.append(np.mean(pred_L2_err_test_ave_cubic))\n",
    "    pred_L2_err_test_adap_std_list2_cubic.append(np.std(pred_L2_err_test_ave_cubic))\n",
    "    \n",
    "    print('\\nL2 predicting error of testing data (cubic): mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave_cubic),\n",
    "                                                                         np.std(pred_L2_err_test_ave_cubic)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1a42815",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# compute errors\n",
    "\n",
    "pred_L2_err_test_adap_list_mean_mat1 = np.zeros((num_rep, len(num_grids_adap_list1)))\n",
    "\n",
    "for i in range(len(num_grids_adap_list1)):\n",
    "    pred_L2_err_test_adap_list_mean_mat1[:,i] = torch.stack(pred_L2_err_test_adap_list1_cubic[i]).mean(axis=(1,)).numpy()\n",
    "\n",
    "    \n",
    "pred_L2_err_test_adap_list_mean_mat2 = np.zeros((num_rep, len(num_grids_adap_list2)))\n",
    "\n",
    "for i in range(len(num_grids_adap_list2)):\n",
    "    pred_L2_err_test_adap_list_mean_mat2[:,i] = torch.stack(pred_L2_err_test_adap_list2_cubic[i]).mean(axis=(1,)).numpy()\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d120437b",
   "metadata": {},
   "source": [
    "# RNN models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f21ffaf7",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# RNN, number of grids >= 16\n",
    "\n",
    "num_grids_RNN_list1 = np.arange(16,68,4, dtype=int)+1#np.arange(24,102,6, dtype=int)+1\n",
    "\n",
    "rnnfunc_RNN_list1 = []\n",
    "fit_window_train_RNN_list1, fit_window_test_RNN_list1 = [], []\n",
    "fit_L2_err_train_RNN_list1, fit_L2_err_test_RNN_list1 = [], []\n",
    "fit_L1_err_train_RNN_list1, fit_L1_err_test_RNN_list1 = [], []\n",
    "\n",
    "fit_L2_err_train_RNN_mean_list1, fit_L2_err_test_RNN_mean_list1 = [], []\n",
    "fit_L2_err_train_RNN_std_list1, fit_L2_err_test_RNN_std_list1 = [], []\n",
    "\n",
    "pred_window_train_RNN_list1, pred_window_test_RNN_list1 = [], []\n",
    "pred_window_train_RNN_list1_pre, pred_window_test_RNN_list1_pre = [], []\n",
    "pred_L2_err_train_RNN_list1, pred_L2_err_test_RNN_list1 = [], []\n",
    "pred_L1_err_train_RNN_list1, pred_L1_err_test_RNN_list1 = [], []\n",
    "\n",
    "pred_L2_err_train_RNN_mean_list1, pred_L2_err_test_RNN_mean_list1 = [], []\n",
    "pred_L2_err_train_RNN_std_list1, pred_L2_err_test_RNN_std_list1 = [], []\n",
    "\n",
    "rng = np.random.default_rng() \n",
    "    \n",
    "batch_size = 64#128\n",
    "verbose = 2\n",
    "buffer_start_steps = 2\n",
    "n_iter = 40\n",
    "\n",
    "t1 = time.time()\n",
    "index = np.arange(num_fine_grid+1)\n",
    "for i, num_grids in enumerate(num_grids_RNN_list1):\n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "    print('\\nnon-adaptive method, number of grids = %d\\n' % num_grids)\n",
    "    index1 = np.round(np.arange(num_grids-1) * num_fine_grid/(num_grids-1))\n",
    "    index1 = np.concatenate(([index1[0]]*buffer_start_steps, index1, [num_fine_grid])) \n",
    "\n",
    "    print('train models:\\n')\n",
    "\n",
    "    rnn1_list = []\n",
    "    fit_window_train1_list, fit_window_test1_list = [], []\n",
    "    fit_L2_err_train1_list, fit_L2_err_test1_list = [], []\n",
    "    fit_L1_err_train1_list, fit_L1_err_test1_list = [], []\n",
    "\n",
    "    pred_window_train1_list, pred_window_test1_list = [], []\n",
    "    pred_window_train1_list_pre, pred_window_test1_list_pre = [], []\n",
    "    pred_L2_err_train1_list, pred_L2_err_test1_list = [], []\n",
    "    pred_L1_err_train1_list, pred_L1_err_test1_list = [], []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "                    \n",
    "        train_ts1 = torch.tensor(np.linspace(train_ts0[:,0], train_ts0[:,-1], num_grids).T)\n",
    "        train_ts1 = torch.concat((train_ts1[:,0].reshape(train_ts0.shape[0],1).repeat(1,buffer_start_steps), \n",
    "                                train_ts1),axis=1)\n",
    "        train_windows1 = interpolate_window(train_windows0, train_ts0, train_ts1, 'cubic')\n",
    "        \n",
    "        if buffer_start_steps > 0:\n",
    "            deltat = train_ts1[0][-1] - train_ts1[0][-2]\n",
    "            train_ts1[:,:buffer_start_steps] = (train_ts1[:,:buffer_start_steps] + \n",
    "                                                torch.arange(-deltat*buffer_start_steps, 0, deltat))\n",
    "            \n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "\n",
    "        train_loader1 = torch.utils.data.DataLoader(\n",
    "            MyDataset(train_windows=train_windows1[train_index,:,:], train_ts=train_ts1[train_index,:]), \n",
    "            shuffle=True, batch_size=batch_size,) # rescale the data\n",
    "\n",
    "        input_steps1 = train_windows1.shape[1]  \n",
    "\n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "        flag, rnn1 = train_RNN(train_loader1, input_steps1, valid_windows0, valid_ts0, num_grids, verbose, \n",
    "                                n_iter=n_iter, thres1=4e2, n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale )   \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "        \n",
    "        print('\\nfitting error of training data:')\n",
    "        fit_window_train1, fit_L2_err_train1, fit_L1_err_train1 = fit_RNN_grids(\n",
    "            rnn1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\nfitting error of testing data:')\n",
    "        fit_window_test1, fit_L2_err_test1, fit_L1_err_test1 = fit_RNN_grids(\n",
    "            rnn1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "\n",
    "        rnn1_list.append((copy.deepcopy(rnn1)))\n",
    "        \n",
    "        fit_L2_err_train1_list.append(fit_L2_err_train1)\n",
    "        fit_L2_err_test1_list.append(fit_L2_err_test1)\n",
    "        fit_L1_err_train1_list.append(fit_L1_err_train1)\n",
    "        fit_L1_err_test1_list.append(fit_L1_err_test1)\n",
    "        fit_window_train1_list.append(fit_window_train1)\n",
    "        fit_window_test1_list.append(fit_window_test1)\n",
    "\n",
    "        print('\\npredicting error of training data:')\n",
    "        pred_window_train1_pre, pred_window_train1, pred_L2_err_train1, pred_L1_err_train1 = pred_RNN_grids(\n",
    "            rnn1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\npredicting error of testing data:')\n",
    "        pred_window_test1_pre, pred_window_test1, pred_L2_err_test1, pred_L1_err_test1 = pred_RNN_grids(\n",
    "            rnn1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "        pred_L2_err_train1_list.append(pred_L2_err_train1)\n",
    "        pred_L2_err_test1_list.append(pred_L2_err_test1)\n",
    "        pred_L1_err_train1_list.append(pred_L1_err_train1)\n",
    "        pred_L1_err_test1_list.append(pred_L1_err_test1)\n",
    "        pred_window_train1_list.append(pred_window_train1)\n",
    "        pred_window_test1_list.append(pred_window_test1)\n",
    "        pred_window_train1_list_pre.append(pred_window_train1_pre)\n",
    "        pred_window_test1_list_pre.append(pred_window_test1_pre)\n",
    "\n",
    "        \n",
    "    rnnfunc_RNN_list1.append(rnn1_list)\n",
    "\n",
    "    fit_window_train_RNN_list1.append(fit_window_train1_list)\n",
    "    fit_window_test_RNN_list1.append(fit_window_test1_list)\n",
    "    fit_L2_err_train_RNN_list1.append(fit_L2_err_train1_list)\n",
    "    fit_L2_err_test_RNN_list1.append(fit_L2_err_test1_list)\n",
    "    fit_L1_err_train_RNN_list1.append(fit_L1_err_train1_list)\n",
    "    fit_L1_err_test_RNN_list1.append(fit_L1_err_test1_list)    \n",
    "    \n",
    "    fit_L2_err_train_ave = np.mean(torch.stack(fit_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave = np.mean(torch.stack(fit_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_train_RNN_mean_list1.append(np.mean(fit_L2_err_train_ave))\n",
    "    fit_L2_err_train_RNN_std_list1.append(np.std(fit_L2_err_train_ave))\n",
    "    fit_L2_err_test_RNN_mean_list1.append(np.mean(fit_L2_err_test_ave))\n",
    "    fit_L2_err_test_RNN_std_list1.append(np.std(fit_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave),\n",
    "                                                                         np.std(fit_L2_err_train_ave)))\n",
    "    print('\\nL2 fitting error of testing data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_test_ave),\n",
    "                                                                         np.std(fit_L2_err_test_ave)))\n",
    "    pred_window_train_RNN_list1.append(pred_window_train1_list)\n",
    "    pred_window_test_RNN_list1.append(pred_window_test1_list)\n",
    "    pred_window_train_RNN_list1_pre.append(pred_window_train1_list_pre)\n",
    "    pred_window_test_RNN_list1_pre.append(pred_window_test1_list_pre)\n",
    "    pred_L2_err_train_RNN_list1.append(pred_L2_err_train1_list)\n",
    "    pred_L2_err_test_RNN_list1.append(pred_L2_err_test1_list)\n",
    "    pred_L1_err_train_RNN_list1.append(pred_L1_err_train1_list)\n",
    "    pred_L1_err_test_RNN_list1.append(pred_L1_err_test1_list)    \n",
    "    \n",
    "    pred_L2_err_train_ave = np.mean(torch.stack(pred_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_test_ave = np.mean(torch.stack(pred_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_train_RNN_mean_list1.append(np.mean(pred_L2_err_train_ave))\n",
    "    pred_L2_err_train_RNN_std_list1.append(np.std(pred_L2_err_train_ave))\n",
    "    pred_L2_err_test_RNN_mean_list1.append(np.mean(pred_L2_err_test_ave))\n",
    "    pred_L2_err_test_RNN_std_list1.append(np.std(pred_L2_err_test_ave))\n",
    "\n",
    "    print('\\nL2 predicting error of training data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_train_ave),\n",
    "                                                                         np.std(pred_L2_err_train_ave)))\n",
    "    print('\\nL2 predicting error of testing data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave),\n",
    "                                                                         np.std(pred_L2_err_test_ave)))\n",
    "    \n",
    "t2 = time.time()\n",
    "print('training model takes %.2f s' % (t2-t1) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc30ff94",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# RNN, number of grids < 16\n",
    "\n",
    "\n",
    "num_grids_RNN_list2 = np.array([6, 8, 12])+1#np.arange(24,102,6, dtype=int)+1\n",
    "\n",
    "rnnfunc_RNN_list2 = []\n",
    "fit_window_train_RNN_list2, fit_window_test_RNN_list2 = [], []\n",
    "fit_L2_err_train_RNN_list2, fit_L2_err_test_RNN_list2 = [], []\n",
    "fit_L1_err_train_RNN_list2, fit_L1_err_test_RNN_list2 = [], []\n",
    "\n",
    "fit_L2_err_train_RNN_mean_list2, fit_L2_err_test_RNN_mean_list2 = [], []\n",
    "fit_L2_err_train_RNN_std_list2, fit_L2_err_test_RNN_std_list2 = [], []\n",
    "\n",
    "pred_window_train_RNN_list2, pred_window_test_RNN_list2 = [], []\n",
    "pred_window_train_RNN_list2_pre, pred_window_test_RNN_list2_pre = [], []\n",
    "pred_L2_err_train_RNN_list2, pred_L2_err_test_RNN_list2 = [], []\n",
    "pred_L1_err_train_RNN_list2, pred_L1_err_test_RNN_list2 = [], []\n",
    "\n",
    "pred_L2_err_train_RNN_mean_list2, pred_L2_err_test_RNN_mean_list2 = [], []\n",
    "pred_L2_err_train_RNN_std_list2, pred_L2_err_test_RNN_std_list2 = [], []\n",
    "\n",
    "rng = np.random.default_rng() \n",
    "    \n",
    "batch_size = 64#128\n",
    "verbose = 2\n",
    "buffer_start_steps = 2\n",
    "\n",
    "t1 = time.time()\n",
    "index = np.arange(num_fine_grid+1)\n",
    "for i, num_grids in enumerate(num_grids_RNN_list2):\n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "    print('\\nnon-adaptive method, number of grids = %d\\n' % num_grids)\n",
    "    index1 = np.round(np.arange(num_grids-1) * num_fine_grid/(num_grids-1))\n",
    "    index1 = np.concatenate(([index1[0]]*buffer_start_steps, index1, [num_fine_grid])) \n",
    "\n",
    "    print('train models:\\n')\n",
    "\n",
    "    rnn1_list = []\n",
    "    fit_window_train1_list, fit_window_test1_list = [], []\n",
    "    fit_L2_err_train1_list, fit_L2_err_test1_list = [], []\n",
    "    fit_L1_err_train1_list, fit_L1_err_test1_list = [], []\n",
    "\n",
    "    pred_window_train1_list, pred_window_test1_list = [], []\n",
    "    pred_window_train1_list_pre, pred_window_test1_list_pre = [], []\n",
    "    pred_L2_err_train1_list, pred_L2_err_test1_list = [], []\n",
    "    pred_L1_err_train1_list, pred_L1_err_test1_list = [], []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "                    \n",
    "        train_ts1 = torch.tensor(np.linspace(train_ts0[:,0], train_ts0[:,-1], num_grids).T)\n",
    "        train_ts1 = torch.concat((train_ts1[:,0].reshape(train_ts0.shape[0],1).repeat(1,buffer_start_steps), \n",
    "                                train_ts1),axis=1)\n",
    "        train_windows1 = interpolate_window(train_windows0, train_ts0, train_ts1, 'cubic')\n",
    "        \n",
    "        if buffer_start_steps > 0:\n",
    "            deltat = train_ts1[0][-1] - train_ts1[0][-2]\n",
    "            train_ts1[:,:buffer_start_steps] = (train_ts1[:,:buffer_start_steps] + \n",
    "                                                torch.arange(-deltat*buffer_start_steps, 0, deltat))\n",
    "            \n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "\n",
    "        train_loader1 = torch.utils.data.DataLoader(\n",
    "            MyDataset(train_windows=train_windows1[train_index,:,:], train_ts=train_ts1[train_index,:]), \n",
    "            shuffle=True, batch_size=batch_size) # rescale the data\n",
    "\n",
    "        input_steps1 = train_windows1.shape[1]  \n",
    "\n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "        flag, rnn1 = train_RNN(train_loader1, input_steps1, valid_windows0, valid_ts0, num_grids, verbose, \n",
    "                                n_iter=n_iter, thres1=4e2, n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale)    \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "        \n",
    "        print('\\nfitting error of training data:')\n",
    "        fit_window_train1, fit_L2_err_train1, fit_L1_err_train1 = fit_RNN_grids(\n",
    "            rnn1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\nfitting error of testing data:')\n",
    "        fit_window_test1, fit_L2_err_test1, fit_L1_err_test1 = fit_RNN_grids(\n",
    "            rnn1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "\n",
    "        rnn1_list.append((copy.deepcopy(rnn1)))\n",
    "        \n",
    "        fit_L2_err_train1_list.append(fit_L2_err_train1)\n",
    "        fit_L2_err_test1_list.append(fit_L2_err_test1)\n",
    "        fit_L1_err_train1_list.append(fit_L1_err_train1)\n",
    "        fit_L1_err_test1_list.append(fit_L1_err_test1)\n",
    "        fit_window_train1_list.append(fit_window_train1)\n",
    "        fit_window_test1_list.append(fit_window_test1)\n",
    "\n",
    "        print('\\npredicting error of training data:')\n",
    "        pred_window_train1_pre, pred_window_train1, pred_L2_err_train1, pred_L1_err_train1 = pred_RNN_grids(\n",
    "            rnn1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\npredicting error of testing data:')\n",
    "        pred_window_test1_pre, pred_window_test1, pred_L2_err_test1, pred_L1_err_test1 = pred_RNN_grids(\n",
    "            rnn1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "        pred_L2_err_train1_list.append(pred_L2_err_train1)\n",
    "        pred_L2_err_test1_list.append(pred_L2_err_test1)\n",
    "        pred_L1_err_train1_list.append(pred_L1_err_train1)\n",
    "        pred_L1_err_test1_list.append(pred_L1_err_test1)\n",
    "        pred_window_train1_list.append(pred_window_train1)\n",
    "        pred_window_test1_list.append(pred_window_test1)\n",
    "        pred_window_train1_list_pre.append(pred_window_train1_pre)\n",
    "        pred_window_test1_list_pre.append(pred_window_test1_pre)\n",
    "\n",
    "        \n",
    "    rnnfunc_RNN_list2.append(rnn1_list)\n",
    "\n",
    "    fit_window_train_RNN_list2.append(fit_window_train1_list)\n",
    "    fit_window_test_RNN_list2.append(fit_window_test1_list)\n",
    "    fit_L2_err_train_RNN_list2.append(fit_L2_err_train1_list)\n",
    "    fit_L2_err_test_RNN_list2.append(fit_L2_err_test1_list)\n",
    "    fit_L1_err_train_RNN_list2.append(fit_L1_err_train1_list)\n",
    "    fit_L1_err_test_RNN_list2.append(fit_L1_err_test1_list)    \n",
    "    \n",
    "    fit_L2_err_train_ave = np.mean(torch.stack(fit_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave = np.mean(torch.stack(fit_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_train_RNN_mean_list2.append(np.mean(fit_L2_err_train_ave))\n",
    "    fit_L2_err_train_RNN_std_list2.append(np.std(fit_L2_err_train_ave))\n",
    "    fit_L2_err_test_RNN_mean_list2.append(np.mean(fit_L2_err_test_ave))\n",
    "    fit_L2_err_test_RNN_std_list2.append(np.std(fit_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave),\n",
    "                                                                         np.std(fit_L2_err_train_ave)))\n",
    "    print('\\nL2 fitting error of testing data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_test_ave),\n",
    "                                                                         np.std(fit_L2_err_test_ave)))\n",
    "    pred_window_train_RNN_list2.append(pred_window_train1_list)\n",
    "    pred_window_test_RNN_list2.append(pred_window_test1_list)\n",
    "    pred_window_train_RNN_list2_pre.append(pred_window_train1_list_pre)\n",
    "    pred_window_test_RNN_list2_pre.append(pred_window_test1_list_pre)\n",
    "    pred_L2_err_train_RNN_list2.append(pred_L2_err_train1_list)\n",
    "    pred_L2_err_test_RNN_list2.append(pred_L2_err_test1_list)\n",
    "    pred_L1_err_train_RNN_list2.append(pred_L1_err_train1_list)\n",
    "    pred_L1_err_test_RNN_list2.append(pred_L1_err_test1_list)    \n",
    "    \n",
    "    pred_L2_err_train_ave = np.mean(torch.stack(pred_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_test_ave = np.mean(torch.stack(pred_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_train_RNN_mean_list2.append(np.mean(pred_L2_err_train_ave))\n",
    "    pred_L2_err_train_RNN_std_list2.append(np.std(pred_L2_err_train_ave))\n",
    "    pred_L2_err_test_RNN_mean_list2.append(np.mean(pred_L2_err_test_ave))\n",
    "    pred_L2_err_test_RNN_std_list2.append(np.std(pred_L2_err_test_ave))\n",
    "\n",
    "    print('\\nL2 predicting error of training data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_train_ave),\n",
    "                                                                         np.std(pred_L2_err_train_ave)))\n",
    "    print('\\nL2 predicting error of testing data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave),\n",
    "                                                                         np.std(pred_L2_err_test_ave)))\n",
    "    \n",
    "t2 = time.time()\n",
    "print('training model takes %.2f s' % (t2-t1) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db3768a",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# compute errors\n",
    "\n",
    "pred_L2_err_test_RNN_list_mean_mat1 = np.zeros((num_rep, len(num_grids_RNN_list1)))\n",
    "\n",
    "for i in range(len(num_grids_RNN_list1)):\n",
    "    pred_L2_err_test_RNN_list_mean_mat1[:,i] = torch.stack(pred_L2_err_test_RNN_list1[i]).mean(axis=(1,)).numpy()\n",
    "    \n",
    "pred_L2_err_test_RNN_list_mean_mat2 = np.zeros((num_rep, len(num_grids_RNN_list2)))\n",
    "\n",
    "for i in range(len(num_grids_RNN_list2)):\n",
    "    pred_L2_err_test_RNN_list_mean_mat2[:,i] = torch.stack(pred_L2_err_test_RNN_list2[i]).mean(axis=(1,)).numpy()\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a17563c1",
   "metadata": {},
   "source": [
    "# LSTM models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ac1bc00",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# LSTM, number of grids >= 16\n",
    "\n",
    "num_grids_LSTM_list1 = np.arange(16,68,4, dtype=int)+1#np.arange(24,102,6, dtype=int)+1\n",
    "\n",
    "lstmfunc_LSTM_list1 = []\n",
    "fit_window_train_LSTM_list1, fit_window_test_LSTM_list1 = [], []\n",
    "fit_L2_err_train_LSTM_list1, fit_L2_err_test_LSTM_list1 = [], []\n",
    "fit_L1_err_train_LSTM_list1, fit_L1_err_test_LSTM_list1 = [], []\n",
    "\n",
    "fit_L2_err_train_LSTM_mean_list1, fit_L2_err_test_LSTM_mean_list1 = [], []\n",
    "fit_L2_err_train_LSTM_std_list1, fit_L2_err_test_LSTM_std_list1 = [], []\n",
    "\n",
    "pred_window_train_LSTM_list1, pred_window_test_LSTM_list1 = [], []\n",
    "pred_window_train_LSTM_list1_pre, pred_window_test_LSTM_list1_pre = [], []\n",
    "pred_L2_err_train_LSTM_list1, pred_L2_err_test_LSTM_list1 = [], []\n",
    "pred_L1_err_train_LSTM_list1, pred_L1_err_test_LSTM_list1 = [], []\n",
    "\n",
    "pred_L2_err_train_LSTM_mean_list1, pred_L2_err_test_LSTM_mean_list1 = [], []\n",
    "pred_L2_err_train_LSTM_std_list1, pred_L2_err_test_LSTM_std_list1 = [], []\n",
    "\n",
    "rng = np.random.default_rng() \n",
    "    \n",
    "\n",
    "batch_size = 64#128\n",
    "verbose = 2\n",
    "\n",
    "\n",
    "t1 = time.time()\n",
    "index = np.arange(num_fine_grid+1)\n",
    "for i, num_grids in enumerate(num_grids_LSTM_list1):\n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "    print('\\nnon-adaptive method, number of grids = %d\\n' % num_grids)\n",
    "    index1 = np.round(np.arange(num_grids-1) * num_fine_grid/(num_grids-1))\n",
    "    index1 = np.concatenate(([index1[0]]*buffer_start_steps, index1, [num_fine_grid])) \n",
    "\n",
    "    print('train models:\\n')\n",
    "\n",
    "    LSTM1_list = []\n",
    "    fit_window_train1_list, fit_window_test1_list = [], []\n",
    "    fit_L2_err_train1_list, fit_L2_err_test1_list = [], []\n",
    "    fit_L1_err_train1_list, fit_L1_err_test1_list = [], []\n",
    "\n",
    "    pred_window_train1_list, pred_window_test1_list = [], []\n",
    "    pred_window_train1_list_pre, pred_window_test1_list_pre = [], []\n",
    "    pred_L2_err_train1_list, pred_L2_err_test1_list = [], []\n",
    "    pred_L1_err_train1_list, pred_L1_err_test1_list = [], []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "                    \n",
    "        train_ts1 = torch.tensor(np.linspace(train_ts0[:,0], train_ts0[:,-1], num_grids).T)\n",
    "        train_ts1 = torch.concat((train_ts1[:,0].reshape(train_ts0.shape[0],1).repeat(1,buffer_start_steps), \n",
    "                                train_ts1),axis=1)\n",
    "        train_windows1 = interpolate_window(train_windows0, train_ts0, train_ts1, 'cubic')\n",
    "        \n",
    "        if buffer_start_steps > 0:\n",
    "            deltat = train_ts1[0][-1] - train_ts1[0][-2]\n",
    "            train_ts1[:,:buffer_start_steps] = (train_ts1[:,:buffer_start_steps] + \n",
    "                                                torch.arange(-deltat*buffer_start_steps, 0, deltat))\n",
    "            \n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "\n",
    "        train_loader1 = torch.utils.data.DataLoader(\n",
    "            MyDataset(train_windows=train_windows1[train_index,:,:], train_ts=train_ts1[train_index,:]), \n",
    "            shuffle=True, batch_size=batch_size,) # rescale the data\n",
    "\n",
    "        input_steps1 = train_windows1.shape[1]  \n",
    "\n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "        flag, lstm1 = train_LSTM(train_loader1, input_steps1, valid_windows0, valid_ts0, num_grids, \n",
    "                                  verbose, n_iter=n_iter, thres1=6e2, thres2=5e2,\n",
    "                                  n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale)    \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "        \n",
    "        print('\\nfitting error of training data:')\n",
    "        fit_window_train1, fit_L2_err_train1, fit_L1_err_train1 = fit_LSTM_grids(\n",
    "            lstm1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\nfitting error of testing data:')\n",
    "        fit_window_test1, fit_L2_err_test1, fit_L1_err_test1 = fit_LSTM_grids(\n",
    "            lstm1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "\n",
    "        LSTM1_list.append(copy.deepcopy(lstm1))\n",
    "        \n",
    "        fit_L2_err_train1_list.append(fit_L2_err_train1)\n",
    "        fit_L2_err_test1_list.append(fit_L2_err_test1)\n",
    "        fit_L1_err_train1_list.append(fit_L1_err_train1)\n",
    "        fit_L1_err_test1_list.append(fit_L1_err_test1)\n",
    "        fit_window_train1_list.append(fit_window_train1)\n",
    "        fit_window_test1_list.append(fit_window_test1)\n",
    "\n",
    "        print('\\npredicting error of training data:')\n",
    "        pred_window_train1_pre, pred_window_train1, pred_L2_err_train1, pred_L1_err_train1 = pred_LSTM_grids(\n",
    "            lstm1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\npredicting error of testing data:')\n",
    "        pred_window_test1_pre, pred_window_test1, pred_L2_err_test1, pred_L1_err_test1 = pred_LSTM_grids(\n",
    "            lstm1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "        pred_L2_err_train1_list.append(pred_L2_err_train1)\n",
    "        pred_L2_err_test1_list.append(pred_L2_err_test1)\n",
    "        pred_L1_err_train1_list.append(pred_L1_err_train1)\n",
    "        pred_L1_err_test1_list.append(pred_L1_err_test1)\n",
    "        pred_window_train1_list.append(pred_window_train1)\n",
    "        pred_window_test1_list.append(pred_window_test1)\n",
    "        pred_window_train1_list_pre.append(pred_window_train1_pre)\n",
    "        pred_window_test1_list_pre.append(pred_window_test1_pre)\n",
    "\n",
    "    lstmfunc_LSTM_list1.append(LSTM1_list)\n",
    "\n",
    "    fit_window_train_LSTM_list1.append(fit_window_train1_list)\n",
    "    fit_window_test_LSTM_list1.append(fit_window_test1_list)\n",
    "    fit_L2_err_train_LSTM_list1.append(fit_L2_err_train1_list)\n",
    "    fit_L2_err_test_LSTM_list1.append(fit_L2_err_test1_list)\n",
    "    fit_L1_err_train_LSTM_list1.append(fit_L1_err_train1_list)\n",
    "    fit_L1_err_test_LSTM_list1.append(fit_L1_err_test1_list)    \n",
    "    \n",
    "    fit_L2_err_train_ave = np.mean(torch.stack(fit_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave = np.mean(torch.stack(fit_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_train_LSTM_mean_list1.append(np.mean(fit_L2_err_train_ave))\n",
    "    fit_L2_err_train_LSTM_std_list1.append(np.std(fit_L2_err_train_ave))\n",
    "    fit_L2_err_test_LSTM_mean_list1.append(np.mean(fit_L2_err_test_ave))\n",
    "    fit_L2_err_test_LSTM_std_list1.append(np.std(fit_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave),\n",
    "                                                                         np.std(fit_L2_err_train_ave)))\n",
    "    print('\\nL2 fitting error of testing data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_test_ave),\n",
    "                                                                         np.std(fit_L2_err_test_ave)))\n",
    "\n",
    "    pred_window_train_LSTM_list1.append(pred_window_train1_list)\n",
    "    pred_window_test_LSTM_list1.append(pred_window_test1_list)\n",
    "    pred_window_train_LSTM_list1_pre.append(pred_window_train1_list_pre)\n",
    "    pred_window_test_LSTM_list1_pre.append(pred_window_test1_list_pre)\n",
    "    pred_L2_err_train_LSTM_list1.append(pred_L2_err_train1_list)\n",
    "    pred_L2_err_test_LSTM_list1.append(pred_L2_err_test1_list)\n",
    "    pred_L1_err_train_LSTM_list1.append(pred_L1_err_train1_list)\n",
    "    pred_L1_err_test_LSTM_list1.append(pred_L1_err_test1_list)    \n",
    "    \n",
    "    pred_L2_err_train_ave = np.mean(torch.stack(pred_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_test_ave = np.mean(torch.stack(pred_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_train_LSTM_mean_list1.append(np.mean(pred_L2_err_train_ave))\n",
    "    pred_L2_err_train_LSTM_std_list1.append(np.std(pred_L2_err_train_ave))\n",
    "    pred_L2_err_test_LSTM_mean_list1.append(np.mean(pred_L2_err_test_ave))\n",
    "    pred_L2_err_test_LSTM_std_list1.append(np.std(pred_L2_err_test_ave))\n",
    "\n",
    "    print('\\nL2 predicting error of training data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_train_ave),\n",
    "                                                                         np.std(pred_L2_err_train_ave)))\n",
    "    print('\\nL2 predicting error of testing data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave),\n",
    "                                                                         np.std(pred_L2_err_test_ave)))\n",
    "\n",
    "t2 = time.time()\n",
    "print('training model takes %.2f s' % (t2-t1) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94bc2d02",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# LSTM, number of grids < 16\n",
    "\n",
    "num_grids_LSTM_list2 = np.array([6,8,12])+1#np.arange(24,102,6, dtype=int)+1\n",
    "\n",
    "lstmfunc_LSTM_list2 = []\n",
    "fit_window_train_LSTM_list2, fit_window_test_LSTM_list2 = [], []\n",
    "fit_L2_err_train_LSTM_list2, fit_L2_err_test_LSTM_list2 = [], []\n",
    "fit_L1_err_train_LSTM_list2, fit_L1_err_test_LSTM_list2 = [], []\n",
    "\n",
    "fit_L2_err_train_LSTM_mean_list2, fit_L2_err_test_LSTM_mean_list2 = [], []\n",
    "fit_L2_err_train_LSTM_std_list2, fit_L2_err_test_LSTM_std_list2 = [], []\n",
    "\n",
    "pred_window_train_LSTM_list2, pred_window_test_LSTM_list2 = [], []\n",
    "pred_window_train_LSTM_list2_pre, pred_window_test_LSTM_list2_pre = [], []\n",
    "pred_L2_err_train_LSTM_list2, pred_L2_err_test_LSTM_list2 = [], []\n",
    "pred_L1_err_train_LSTM_list2, pred_L1_err_test_LSTM_list2 = [], []\n",
    "\n",
    "pred_L2_err_train_LSTM_mean_list2, pred_L2_err_test_LSTM_mean_list2 = [], []\n",
    "pred_L2_err_train_LSTM_std_list2, pred_L2_err_test_LSTM_std_list2 = [], []\n",
    "\n",
    "rng = np.random.default_rng() \n",
    "    \n",
    "\n",
    "batch_size = 64#128\n",
    "verbose = 2\n",
    "\n",
    "\n",
    "t1 = time.time()\n",
    "index = np.arange(num_fine_grid+1)\n",
    "for i, num_grids in enumerate(num_grids_LSTM_list2):\n",
    "    print('\\n----------------------------------------------------------------\\n')\n",
    "    print('\\nnon-adaptive method, number of grids = %d\\n' % num_grids)\n",
    "    index1 = np.round(np.arange(num_grids-1) * num_fine_grid/(num_grids-1))\n",
    "    index1 = np.concatenate(([index1[0]]*buffer_start_steps, index1, [num_fine_grid])) \n",
    "\n",
    "    print('train models:\\n')\n",
    "\n",
    "    LSTM1_list = []\n",
    "    fit_window_train1_list, fit_window_test1_list = [], []\n",
    "    fit_L2_err_train1_list, fit_L2_err_test1_list = [], []\n",
    "    fit_L1_err_train1_list, fit_L1_err_test1_list = [], []\n",
    "\n",
    "    pred_window_train1_list, pred_window_test1_list = [], []\n",
    "    pred_window_train1_list_pre, pred_window_test1_list_pre = [], []\n",
    "    pred_L2_err_train1_list, pred_L2_err_test1_list = [], []\n",
    "    pred_L1_err_train1_list, pred_L1_err_test1_list = [], []\n",
    "\n",
    "    rep = 0\n",
    "    while rep <= num_rep-1:\n",
    "        \n",
    "        train_windows0 = train_windows0_list[rep]\n",
    "        train_ts0 = train_ts0_list[rep]\n",
    "        test_windows0 = test_windows0_list[rep]\n",
    "        test_ts0 = test_ts0_list[rep]\n",
    "                    \n",
    "        train_ts1 = torch.tensor(np.linspace(train_ts0[:,0], train_ts0[:,-1], num_grids).T)\n",
    "        train_ts1 = torch.concat((train_ts1[:,0].reshape(train_ts0.shape[0],1).repeat(1,buffer_start_steps), \n",
    "                                train_ts1),axis=1)\n",
    "        train_windows1 = interpolate_window(train_windows0, train_ts0, train_ts1, 'cubic')\n",
    "        \n",
    "        if buffer_start_steps > 0:\n",
    "            deltat = train_ts1[0][-1] - train_ts1[0][-2]\n",
    "            train_ts1[:,:buffer_start_steps] = (train_ts1[:,:buffer_start_steps] + \n",
    "                                                torch.arange(-deltat*buffer_start_steps, 0, deltat))\n",
    "            \n",
    "        valid_index = np.sort(npr.choice(train_windows0.shape[0], 50, replace=False))\n",
    "        valid_windows0 = train_windows0[valid_index,:,:]\n",
    "        valid_ts0 = train_ts0[valid_index,:]\n",
    "        \n",
    "        train_index = np.delete(np.arange(train_windows0.shape[0]), valid_index)     \n",
    "\n",
    "        train_loader1 = torch.utils.data.DataLoader(\n",
    "            MyDataset(train_windows=train_windows1[train_index,:,:], train_ts=train_ts1[train_index,:]), \n",
    "            shuffle=True, batch_size=batch_size,) # rescale the data\n",
    "\n",
    "        input_steps1 = train_windows1.shape[1]  \n",
    "\n",
    "        print('\\n------the %d-th replica:--------' % (rep+1))\n",
    "        flag, lstm1 = train_LSTM(train_loader1, input_steps1, valid_windows0, valid_ts0, num_grids, \n",
    "                                verbose, n_iter=n_iter, thres1=6e2, thres2=5e2,\n",
    "                                n_latent=n_latent, obs_dim=obs_dim, time_scale=time_scale)    \n",
    "        if not flag:\n",
    "            continue\n",
    "        else:\n",
    "            rep += 1\n",
    "        \n",
    "        print('\\nfitting error of training data:')\n",
    "        fit_window_train1, fit_L2_err_train1, fit_L1_err_train1 = fit_LSTM_grids(\n",
    "            lstm1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\nfitting error of testing data:')\n",
    "        fit_window_test1, fit_L2_err_test1, fit_L1_err_test1 = fit_LSTM_grids(\n",
    "            lstm1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "\n",
    "        LSTM1_list.append(copy.deepcopy(lstm1))\n",
    "        \n",
    "        fit_L2_err_train1_list.append(fit_L2_err_train1)\n",
    "        fit_L2_err_test1_list.append(fit_L2_err_test1)\n",
    "        fit_L1_err_train1_list.append(fit_L1_err_train1)\n",
    "        fit_L1_err_test1_list.append(fit_L1_err_test1)\n",
    "        fit_window_train1_list.append(fit_window_train1)\n",
    "        fit_window_test1_list.append(fit_window_test1)\n",
    "\n",
    "        print('\\npredicting error of training data:')\n",
    "        pred_window_train1_pre, pred_window_train1, pred_L2_err_train1, pred_L1_err_train1 = pred_LSTM_grids(\n",
    "            lstm1, train_windows0, train_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "        print('\\npredicting error of testing data:')\n",
    "        pred_window_test1_pre, pred_window_test1, pred_L2_err_test1, pred_L1_err_test1 = pred_LSTM_grids(\n",
    "            lstm1, test_windows0, test_ts0, num_grids, buffer_start_steps=buffer_start_steps,\n",
    "            n_latent=n_latent, obs_dim=obs_dim)\n",
    "\n",
    "        pred_L2_err_train1_list.append(pred_L2_err_train1)\n",
    "        pred_L2_err_test1_list.append(pred_L2_err_test1)\n",
    "        pred_L1_err_train1_list.append(pred_L1_err_train1)\n",
    "        pred_L1_err_test1_list.append(pred_L1_err_test1)\n",
    "        pred_window_train1_list.append(pred_window_train1)\n",
    "        pred_window_test1_list.append(pred_window_test1)\n",
    "        pred_window_train1_list_pre.append(pred_window_train1_pre)\n",
    "        pred_window_test1_list_pre.append(pred_window_test1_pre)\n",
    "\n",
    "    lstmfunc_LSTM_list2.append(LSTM1_list)\n",
    "\n",
    "    fit_window_train_LSTM_list2.append(fit_window_train1_list)\n",
    "    fit_window_test_LSTM_list2.append(fit_window_test1_list)\n",
    "    fit_L2_err_train_LSTM_list2.append(fit_L2_err_train1_list)\n",
    "    fit_L2_err_test_LSTM_list2.append(fit_L2_err_test1_list)\n",
    "    fit_L1_err_train_LSTM_list2.append(fit_L1_err_train1_list)\n",
    "    fit_L1_err_test_LSTM_list2.append(fit_L1_err_test1_list)    \n",
    "    \n",
    "    fit_L2_err_train_ave = np.mean(torch.stack(fit_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_test_ave = np.mean(torch.stack(fit_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    fit_L2_err_train_LSTM_mean_list2.append(np.mean(fit_L2_err_train_ave))\n",
    "    fit_L2_err_train_LSTM_std_list2.append(np.std(fit_L2_err_train_ave))\n",
    "    fit_L2_err_test_LSTM_mean_list2.append(np.mean(fit_L2_err_test_ave))\n",
    "    fit_L2_err_test_LSTM_std_list2.append(np.std(fit_L2_err_test_ave))\n",
    "    \n",
    "    print('\\nL2 fitting error of training data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_train_ave),\n",
    "                                                                         np.std(fit_L2_err_train_ave)))\n",
    "    print('\\nL2 fitting error of testing data: mean = %.2e, std = %.2e' % (np.mean(fit_L2_err_test_ave),\n",
    "                                                                         np.std(fit_L2_err_test_ave)))\n",
    "\n",
    "    pred_window_train_LSTM_list2.append(pred_window_train1_list)\n",
    "    pred_window_test_LSTM_list2.append(pred_window_test1_list)\n",
    "    pred_window_train_LSTM_list2_pre.append(pred_window_train1_list_pre)\n",
    "    pred_window_test_LSTM_list2_pre.append(pred_window_test1_list_pre)\n",
    "    pred_L2_err_train_LSTM_list2.append(pred_L2_err_train1_list)\n",
    "    pred_L2_err_test_LSTM_list2.append(pred_L2_err_test1_list)\n",
    "    pred_L1_err_train_LSTM_list2.append(pred_L1_err_train1_list)\n",
    "    pred_L1_err_test_LSTM_list2.append(pred_L1_err_test1_list)    \n",
    "    \n",
    "    pred_L2_err_train_ave = np.mean(torch.stack(pred_L2_err_train1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_test_ave = np.mean(torch.stack(pred_L2_err_test1_list).numpy(),axis=(1,))\n",
    "    pred_L2_err_train_LSTM_mean_list2.append(np.mean(pred_L2_err_train_ave))\n",
    "    pred_L2_err_train_LSTM_std_list2.append(np.std(pred_L2_err_train_ave))\n",
    "    pred_L2_err_test_LSTM_mean_list2.append(np.mean(pred_L2_err_test_ave))\n",
    "    pred_L2_err_test_LSTM_std_list2.append(np.std(pred_L2_err_test_ave))\n",
    "\n",
    "    print('\\nL2 predicting error of training data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_train_ave),\n",
    "                                                                         np.std(pred_L2_err_train_ave)))\n",
    "    print('\\nL2 predicting error of testing data: mean = %.2e, std = %.2e' % (np.mean(pred_L2_err_test_ave),\n",
    "                                                                         np.std(pred_L2_err_test_ave)))\n",
    "\n",
    "t2 = time.time()\n",
    "print('training model takes %.2f s' % (t2-t1) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e97f4225",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# compute errors\n",
    "\n",
    "pred_L2_err_test_LSTM_list_mean_mat1 = np.zeros((num_rep, len(num_grids_LSTM_list1)))\n",
    "\n",
    "for i in range(len(num_grids_LSTM_list1)):\n",
    "    pred_L2_err_test_LSTM_list_mean_mat1[:,i] = torch.stack(pred_L2_err_test_LSTM_list1[i]).mean(axis=(1,)).numpy()\n",
    "\n",
    "\n",
    "pred_L2_err_test_LSTM_list_mean_mat2 = np.zeros((num_rep, len(num_grids_LSTM_list2)))\n",
    "\n",
    "for i in range(len(num_grids_LSTM_list2)):\n",
    "    pred_L2_err_test_LSTM_list_mean_mat2[:,i] = torch.stack(pred_L2_err_test_LSTM_list2[i]).mean(axis=(1,)).numpy()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac29ad0e",
   "metadata": {},
   "source": [
    "# Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d85705",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# box plot\n",
    "\n",
    "alpha = 0.5\n",
    "index1 = [0,1,2,3,4,5,6,7,8,9,10,11,12]\n",
    "# plt.figure(figsize=(7,5))\n",
    "fig, ax = plt.subplots(figsize=(8,6))\n",
    "linewidth = 2\n",
    "\n",
    "cols = ['tab:green'] * 13\n",
    "xticks = np.concatenate((np.array(num_grids_non_adap_list2), np.array(num_grids_non_adap_list1)))\n",
    "err_temp = np.concatenate((np.array(pred_L2_err_test_non_adap_list_mean_mat2), np.array(pred_L2_err_test_non_adap_list_mean_mat1)), axis=1)\n",
    "bp = ax.boxplot(err_temp, positions=xticks,\n",
    "           widths = 0.4, whis=0.5, whiskerprops={'color' : 'tab:green'}, \n",
    "                medianprops={'color' : 'tab:green', 'linewidth': 3},\n",
    "                boxprops={'color' : 'tab:green'},\n",
    "                flierprops={'markersize':0.01})\n",
    "res = {key : [v.get_ydata() for v in value] for key, value in bp.items()}\n",
    "ax.plot(xticks, np.array(res['medians'])[:,0], 'x-', color='tab:green',\n",
    "        label='RNNODE', alpha=alpha, linewidth=linewidth)\n",
    "for b, c in zip(bp['boxes'], cols):\n",
    "    b.set_alpha(alpha)\n",
    "print(np.array(res['medians'])[:,0])\n",
    "\n",
    "\n",
    "cols = ['tab:red'] * 13\n",
    "xticks = np.concatenate((np.array(num_grids_non_adap_list2), np.array(num_grids_non_adap_list1)))\n",
    "err_temp = np.concatenate((np.array(pred_L2_err_test_adap_list_mean_mat2), np.array(pred_L2_err_test_adap_list_mean_mat1)), axis=1)\n",
    "bp = ax.boxplot(err_temp, positions=xticks,\n",
    "           widths = 0.4, whis=0.5, whiskerprops={'color' : 'tab:red'}, \n",
    "                medianprops={'color' : 'tab:red', 'linewidth': 3},\n",
    "                boxprops={'color' : 'tab:red'},\n",
    "                flierprops={'markersize':0.01})\n",
    "res = {key : [v.get_ydata() for v in value] for key, value in bp.items()}\n",
    "ax.plot(xticks, np.array(res['medians'])[:,0], 'x-', color='tab:red',\n",
    "        label='RNN-ODE-Adap', alpha=alpha, linewidth=linewidth)\n",
    "for b, c in zip(bp['boxes'], cols):\n",
    "    b.set_alpha(alpha)\n",
    "\n",
    "\n",
    "cols = ['tab:blue'] * 13\n",
    "xticks = np.concatenate((np.array(num_grids_RNN_list2), np.array(num_grids_RNN_list1)))\n",
    "err_temp = np.concatenate((np.array(pred_L2_err_test_RNN_list_mean_mat2), np.array(pred_L2_err_test_RNN_list_mean_mat1)), axis=1)\n",
    "bp = ax.boxplot(err_temp, positions=xticks,\n",
    "                 widths = 0.4, whis=0.01, whiskerprops={'color' : 'tab:blue'}, \n",
    "                 medianprops={'color' : 'tab:blue', 'linewidth': 3},\n",
    "                 boxprops={'color' : 'tab:blue'},\n",
    "                 flierprops={'markersize':0.01}\n",
    "                )\n",
    "res = {key : [v.get_ydata() for v in value] for key, value in bp.items()}\n",
    "ax.plot(xticks, np.array(res['medians'])[:,0], 'x-', color='tab:blue',\n",
    "        label='RNN', alpha=alpha, linewidth=linewidth)\n",
    "for b, c in zip(bp['boxes'], cols):\n",
    "    b.set_alpha(alpha)\n",
    "print(np.array(res['medians'])[:,0])\n",
    "\n",
    "cols = ['tab:orange'] * 13\n",
    "xticks = np.concatenate((np.array(num_grids_LSTM_list2), np.array(num_grids_LSTM_list1)))\n",
    "err_temp = np.concatenate((np.array(pred_L2_err_test_LSTM_list_mean_mat2), np.array(pred_L2_err_test_LSTM_list_mean_mat1)), axis=1)\n",
    "bp = ax.boxplot(err_temp, positions=xticks,\n",
    "                 widths = 0.4, whis=0.5, whiskerprops={'color' : 'tab:orange'}, \n",
    "                 medianprops={'color' : 'tab:orange', 'linewidth': 3},\n",
    "                 boxprops={'color' : 'tab:orange'},\n",
    "                 flierprops={'markersize':0.01}\n",
    "                )\n",
    "res = {key : [v.get_ydata() for v in value] for key, value in bp.items()}\n",
    "ax.plot(xticks, np.array(res['medians'])[:,0], 'x-', color='tab:orange',\n",
    "        label='LSTM', alpha=alpha, linewidth=linewidth)\n",
    "for b, c in zip(bp['boxes'], cols):\n",
    "    b.set_alpha(alpha)\n",
    "\n",
    "\n",
    "xticks = np.concatenate((np.array(num_grids_non_adap_list2), np.array(num_grids_non_adap_list1)))\n",
    "plt.xticks(xticks,\n",
    "           list(map(str, xticks)))\n",
    "plt.ylim([0, 0.4])\n",
    "plt.xlabel('number of grids', fontsize=14)\n",
    "plt.ylabel('MSE', fontsize=14)\n",
    "plt.legend(fontsize=14)\n",
    "plt.grid()\n",
    "plt.suptitle('Spiral data: Testing Precition MSE vs. Complexity (boxplot)', fontsize=16)\n",
    "# plt.savefig('spiral_error_boxplot.eps', format='eps', dpi=1000)\n",
    "\n",
    "plt.show()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.9.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}