PyPOTS_Imputation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📑 Tutorials for PyPOTS Imputation Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📀 Preparing the **PhysioNet-2012** dataset for this tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:32:00 [INFO]: Have set the random seed as 16 for numpy and pytorch.\n",
      "2024-05-21 12:32:00 [INFO]: Loading the dataset physionet_2012 with TSDB (https://github.com/WenjieDu/Time_Series_Data_Beans)...\n",
      "2024-05-21 12:32:00 [INFO]: Starting preprocessing physionet_2012...\n",
      "2024-05-21 12:32:00 [INFO]: You're using dataset physionet_2012, please cite it properly in your work. You can find its reference information at the below link: \n",
      "https://github.com/WenjieDu/TSDB/tree/main/dataset_profiles/physionet_2012\n",
      "2024-05-21 12:32:00 [INFO]: Dataset physionet_2012 has already been downloaded. Processing directly...\n",
      "2024-05-21 12:32:00 [INFO]: Dataset physionet_2012 has already been cached. Loading from cache directly...\n",
      "2024-05-21 12:32:00 [INFO]: Loaded successfully!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['n_classes', 'n_steps', 'n_features', 'train_X', 'train_y', 'train_ICUType', 'val_X', 'val_y', 'val_ICUType', 'test_X', 'test_y', 'test_ICUType', 'scaler', 'val_X_ori', 'test_X_ori', 'test_X_indicating_mask'])\n"
     ]
    }
   ],
   "source": [
    "from pypots.data.generating import gene_physionet2012\n",
    "from pypots.utils.random import set_random_seed\n",
    "from global_config import RANDOM_SEED\n",
    "\n",
    "set_random_seed(RANDOM_SEED)\n",
    "\n",
    "# Load the PhysioNet-2012 dataset\n",
    "physionet2012_dataset = gene_physionet2012(artificially_missing_rate=0.1)\n",
    "\n",
    "# Take a look at the generated PhysioNet-2012 dataset, you'll find that everything has been prepared for you,\n",
    "# data splitting, normalization, additional artificially-missing values for evaluation, etc.\n",
    "print(physionet2012_dataset.keys())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the datasets for training, validating, and testing.\n",
    "\n",
    "dataset_for_training = {\n",
    "    \"X\": physionet2012_dataset['train_X'],\n",
    "}\n",
    "\n",
    "dataset_for_validating = {\n",
    "    \"X\": physionet2012_dataset['val_X'],\n",
    "    \"X_ori\": physionet2012_dataset['val_X_ori'],\n",
    "}\n",
    "\n",
    "dataset_for_testing = {\n",
    "    \"X\": physionet2012_dataset['test_X'],\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 An example of **SAITS** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:32:19 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 12:32:19 [INFO]: Model files will be saved to tutorial_results/imputation/saits/20240521_T123219\n",
      "2024-05-21 12:32:19 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/saits/20240521_T123219/tensorboard\n",
      "2024-05-21 12:32:19 [INFO]: SAITS initialized with the given hyperparameters, the number of trainable parameters: 1,378,358\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import SAITS\n",
    "\n",
    "# initialize the model\n",
    "saits = SAITS(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    n_layers=2,\n",
    "    d_model=256,\n",
    "    d_ffn=128,\n",
    "    n_heads=4,\n",
    "    d_k=64,\n",
    "    d_v=64,\n",
    "    dropout=0.1,\n",
    "    attn_dropout=0.1,\n",
    "    diagonal_attention_mask=True,  # otherwise the original self-attention mechanism will be applied\n",
    "    ORT_weight=1,  # you can adjust the weight values of arguments ORT_weight\n",
    "    # and MIT_weight to make the SAITS model focus more on one task. Usually you can just leave them to the default values, i.e. 1.\n",
    "    MIT_weight=1,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,  \n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/saits\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:32:29 [INFO]: Epoch 001 - training loss: 0.7362, validation loss: 0.3747\n",
      "2024-05-21 12:32:36 [INFO]: Epoch 002 - training loss: 0.5380, validation loss: 0.3397\n",
      "2024-05-21 12:32:43 [INFO]: Epoch 003 - training loss: 0.4889, validation loss: 0.3128\n",
      "2024-05-21 12:32:50 [INFO]: Epoch 004 - training loss: 0.4549, validation loss: 0.2992\n",
      "2024-05-21 12:32:57 [INFO]: Epoch 005 - training loss: 0.4289, validation loss: 0.2792\n",
      "2024-05-21 12:33:04 [INFO]: Epoch 006 - training loss: 0.4057, validation loss: 0.2748\n",
      "2024-05-21 12:33:11 [INFO]: Epoch 007 - training loss: 0.3908, validation loss: 0.2692\n",
      "2024-05-21 12:33:18 [INFO]: Epoch 008 - training loss: 0.3803, validation loss: 0.2641\n",
      "2024-05-21 12:33:25 [INFO]: Epoch 009 - training loss: 0.3724, validation loss: 0.2628\n",
      "2024-05-21 12:33:32 [INFO]: Epoch 010 - training loss: 0.3660, validation loss: 0.2613\n",
      "2024-05-21 12:33:32 [INFO]: Finished training. The best model is from epoch#10.\n",
      "2024-05-21 12:33:32 [INFO]: Saved the model to tutorial_results/imputation/saits/20240521_T123219/SAITS.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "saits.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "saits_results = saits.predict(dataset_for_testing)\n",
    "saits_imputation = saits_results[\"imputation\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.2372\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    saits_imputation, \n",
    "    physionet2012_dataset['test_X_ori'], \n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 🚀 An example of **Transformer** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:33:33 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 12:33:33 [INFO]: Model files will be saved to tutorial_results/imputation/transformer/20240521_T123333\n",
      "2024-05-21 12:33:33 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/transformer/20240521_T123333/tensorboard\n",
      "2024-05-21 12:33:33 [INFO]: Transformer initialized with the given hyperparameters, the number of trainable parameters: 7,938,597\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import Transformer\n",
    "\n",
    "# initialize the model\n",
    "transformer = Transformer(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    n_layers=6,\n",
    "    d_model=512,\n",
    "    d_ffn=256,\n",
    "    n_heads=4,\n",
    "    d_k=128,\n",
    "    d_v=128,\n",
    "    dropout=0.1,\n",
    "    attn_dropout=0,\n",
    "    ORT_weight=1,  # you can adjust the weight values of arguments ORT_weight\n",
    "    # and MIT_weight to make the SAITS model focus more on one task. Usually you can just leave them to the default values, i.e. 1.\n",
    "    MIT_weight=1,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,  \n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/transformer\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:33:41 [INFO]: Epoch 001 - training loss: 1.4364, validation loss: 1.0100\n",
      "2024-05-21 12:33:48 [INFO]: Epoch 002 - training loss: 1.3903, validation loss: 0.9911\n",
      "2024-05-21 12:33:55 [INFO]: Epoch 003 - training loss: 1.3806, validation loss: 0.9910\n",
      "2024-05-21 12:34:03 [INFO]: Epoch 004 - training loss: 1.3697, validation loss: 0.9919\n",
      "2024-05-21 12:34:10 [INFO]: Epoch 005 - training loss: 1.3679, validation loss: 0.9864\n",
      "2024-05-21 12:34:18 [INFO]: Epoch 006 - training loss: 1.3655, validation loss: 0.9801\n",
      "2024-05-21 12:34:25 [INFO]: Epoch 007 - training loss: 1.3645, validation loss: 0.9849\n",
      "2024-05-21 12:34:32 [INFO]: Epoch 008 - training loss: 1.3643, validation loss: 0.9886\n",
      "2024-05-21 12:34:40 [INFO]: Epoch 009 - training loss: 1.3642, validation loss: 0.9945\n",
      "2024-05-21 12:34:40 [INFO]: Exceeded the training patience. Terminating the training procedure...\n",
      "2024-05-21 12:34:40 [INFO]: Finished training. The best model is from epoch#6.\n",
      "2024-05-21 12:34:40 [INFO]: Saved the model to tutorial_results/imputation/transformer/20240521_T123333/Transformer.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "transformer.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "transformer_results = transformer.predict(dataset_for_testing)\n",
    "transformer_imputation = transformer_results[\"imputation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.6862\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    transformer_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 An example of **TimesNet** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:34:40 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 12:34:40 [INFO]: Model files will be saved to tutorial_results/imputation/timesnet/20240521_T123440\n",
      "2024-05-21 12:34:40 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/timesnet/20240521_T123440/tensorboard\n",
      "2024-05-21 12:34:41 [INFO]: TimesNet initialized with the given hyperparameters, the number of trainable parameters: 21,649,573\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import TimesNet\n",
    "\n",
    "# initialize the model\n",
    "timesnet = TimesNet(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    n_layers=1,\n",
    "    top_k=1,\n",
    "    d_model=128,\n",
    "    d_ffn=512,\n",
    "    n_kernels=5,\n",
    "    dropout=0.5,\n",
    "    apply_nonstationary_norm=False,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,\n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/timesnet\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:34:48 [INFO]: Epoch 001 - training loss: 0.4796, validation loss: 0.3719\n",
      "2024-05-21 12:34:52 [INFO]: Epoch 002 - training loss: 0.4039, validation loss: 0.3458\n",
      "2024-05-21 12:34:55 [INFO]: Epoch 003 - training loss: 0.4087, validation loss: 0.3382\n",
      "2024-05-21 12:34:59 [INFO]: Epoch 004 - training loss: 0.4472, validation loss: 0.3351\n",
      "2024-05-21 12:35:02 [INFO]: Epoch 005 - training loss: 0.4144, validation loss: 0.3300\n",
      "2024-05-21 12:35:06 [INFO]: Epoch 006 - training loss: 0.4142, validation loss: 0.3274\n",
      "2024-05-21 12:35:09 [INFO]: Epoch 007 - training loss: 0.4203, validation loss: 0.3257\n",
      "2024-05-21 12:35:13 [INFO]: Epoch 008 - training loss: 0.4631, validation loss: 0.3249\n",
      "2024-05-21 12:35:16 [INFO]: Epoch 009 - training loss: 0.3840, validation loss: 0.3172\n",
      "2024-05-21 12:35:20 [INFO]: Epoch 010 - training loss: 0.3685, validation loss: 0.3151\n",
      "2024-05-21 12:35:20 [INFO]: Finished training. The best model is from epoch#10.\n",
      "2024-05-21 12:35:20 [INFO]: Saved the model to tutorial_results/imputation/timesnet/20240521_T123440/TimesNet.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "timesnet.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "timesnet_results = timesnet.predict(dataset_for_testing)\n",
    "timesnet_imputation = timesnet_results[\"imputation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.3243\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    timesnet_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 An example of **CSDI** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:35:21 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 12:35:21 [INFO]: Model files will be saved to tutorial_results/imputation/csdi/20240521_T123521\n",
      "2024-05-21 12:35:21 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/csdi/20240521_T123521/tensorboard\n",
      "/home/linglong/data/linglong/.conda/envs/pypots-dev/lib/python3.11/site-packages/torch/nn/modules/transformer.py:286: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)\n",
      "  warnings.warn(f\"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}\")\n",
      "2024-05-21 12:35:21 [INFO]: CSDI initialized with the given hyperparameters, the number of trainable parameters: 1,694,753\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import CSDI\n",
    "\n",
    "# initialize the model\n",
    "csdi = CSDI(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    n_layers=6,\n",
    "    n_heads=2,\n",
    "    n_channels=128,\n",
    "    d_time_embedding=64,\n",
    "    d_feature_embedding=32,\n",
    "    d_diffusion_embedding=128,\n",
    "    target_strategy=\"random\",\n",
    "    n_diffusion_steps=50,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,\n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/csdi\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 12:37:43 [INFO]: Epoch 001 - training loss: 0.3379, validation loss: 0.2506\n",
      "2024-05-21 12:40:08 [INFO]: Epoch 002 - training loss: 0.2575, validation loss: 0.2083\n",
      "2024-05-21 12:42:31 [INFO]: Epoch 003 - training loss: 0.2406, validation loss: 0.1998\n",
      "2024-05-21 12:44:56 [INFO]: Epoch 004 - training loss: 0.2389, validation loss: 0.1978\n",
      "2024-05-21 12:47:20 [INFO]: Epoch 005 - training loss: 0.2343, validation loss: 0.1984\n",
      "2024-05-21 12:49:44 [INFO]: Epoch 006 - training loss: 0.2272, validation loss: 0.1907\n",
      "2024-05-21 12:52:07 [INFO]: Epoch 007 - training loss: 0.2298, validation loss: 0.1924\n",
      "2024-05-21 12:54:32 [INFO]: Epoch 008 - training loss: 0.2208, validation loss: 0.1882\n",
      "2024-05-21 12:56:56 [INFO]: Epoch 009 - training loss: 0.2267, validation loss: 0.1827\n",
      "2024-05-21 12:59:20 [INFO]: Epoch 010 - training loss: 0.2188, validation loss: 0.1845\n",
      "2024-05-21 12:59:20 [INFO]: Finished training. The best model is from epoch#9.\n",
      "2024-05-21 12:59:20 [INFO]: Saved the model to tutorial_results/imputation/csdi/20240521_T123521/CSDI.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "csdi.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The shape of csdi_imputation is (2398, 2, 48, 37)\n"
     ]
    }
   ],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "\n",
    "# CSDI has an argument to control the number of sampling times during inference\n",
    "csdi_results = csdi.predict(dataset_for_testing, n_sampling_times=2)\n",
    "csdi_imputation = csdi_results[\"imputation\"]\n",
    "\n",
    "print(f\"The shape of csdi_imputation is {csdi_imputation.shape}\")\n",
    "\n",
    "# for error calculation, we need to take the mean value of the multiple samplings for each data sample\n",
    "mean_csdi_imputation = csdi_imputation.mean(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.2781\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    mean_csdi_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 🚀 An example of **US-GAN** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 13:04:08 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 13:04:08 [INFO]: Model files will be saved to tutorial_results/imputation/us_gan/20240521_T130408\n",
      "2024-05-21 13:04:08 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/us_gan/20240521_T130408/tensorboard\n",
      "2024-05-21 13:04:08 [INFO]: USGAN initialized with the given hyperparameters, the number of trainable parameters: 1,258,517\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import USGAN\n",
    "\n",
    "# initialize the model\n",
    "us_gan = USGAN(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    rnn_hidden_size=256,\n",
    "    lambda_mse=1,\n",
    "    dropout=0.1,\n",
    "    G_steps=1,\n",
    "    D_steps=1,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    G_optimizer=Adam(lr=1e-3),\n",
    "    D_optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,\n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/us_gan\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 13:06:21 [INFO]: Epoch 001 - generator training loss: 0.4250, discriminator training loss: 0.1838, validation loss: 0.3584\n",
      "2024-05-21 13:08:16 [INFO]: Epoch 002 - generator training loss: 0.3517, discriminator training loss: 0.0526, validation loss: 0.3129\n",
      "2024-05-21 13:10:10 [INFO]: Epoch 003 - generator training loss: 0.3263, discriminator training loss: 0.0367, validation loss: 0.2942\n",
      "2024-05-21 13:12:00 [INFO]: Epoch 004 - generator training loss: 0.3103, discriminator training loss: 0.0310, validation loss: 0.2862\n",
      "2024-05-21 13:13:53 [INFO]: Epoch 005 - generator training loss: 0.3010, discriminator training loss: 0.0281, validation loss: 0.2839\n",
      "2024-05-21 13:15:47 [INFO]: Epoch 006 - generator training loss: 0.2950, discriminator training loss: 0.0281, validation loss: 0.2732\n",
      "2024-05-21 13:17:40 [INFO]: Epoch 007 - generator training loss: 0.2836, discriminator training loss: 0.0258, validation loss: 0.2791\n",
      "2024-05-21 13:19:37 [INFO]: Epoch 008 - generator training loss: 0.2760, discriminator training loss: 0.0243, validation loss: 0.2653\n",
      "2024-05-21 13:21:33 [INFO]: Epoch 009 - generator training loss: 0.2683, discriminator training loss: 0.0233, validation loss: 0.2699\n",
      "2024-05-21 13:23:29 [INFO]: Epoch 010 - generator training loss: 0.2678, discriminator training loss: 0.0227, validation loss: 0.2701\n",
      "2024-05-21 13:23:29 [INFO]: Finished training. The best model is from epoch#8.\n",
      "2024-05-21 13:23:29 [INFO]: Saved the model to tutorial_results/imputation/us_gan/20240521_T130408/USGAN.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "us_gan.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "us_gan_results = us_gan.predict(dataset_for_testing)\n",
    "us_gan_imputation = us_gan_results[\"imputation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.2805\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    us_gan_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 🚀 An example of **GP-VAE** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 13:23:43 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 13:23:43 [INFO]: Model files will be saved to tutorial_results/imputation/gp_vae/20240521_T132343\n",
      "2024-05-21 13:23:43 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/gp_vae/20240521_T132343/tensorboard\n",
      "2024-05-21 13:23:43 [INFO]: GPVAE initialized with the given hyperparameters, the number of trainable parameters: 229,652\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import GPVAE\n",
    "\n",
    "# initialize the model\n",
    "gp_vae = GPVAE(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    latent_size=37,\n",
    "    encoder_sizes=(128,128),\n",
    "    decoder_sizes=(256,256),\n",
    "    kernel=\"cauchy\",\n",
    "    beta=0.2,\n",
    "    M=1,\n",
    "    K=1,\n",
    "    sigma=1.005,\n",
    "    length_scale=7.0,\n",
    "    kernel_scales=1,\n",
    "    window_size=24,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,\n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/gp_vae\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 13:23:47 [INFO]: Epoch 001 - training loss: 26022.2357, validation loss: 0.6172\n",
      "2024-05-21 13:23:51 [INFO]: Epoch 002 - training loss: 22874.5681, validation loss: 0.5858\n",
      "2024-05-21 13:23:55 [INFO]: Epoch 003 - training loss: 22840.0816, validation loss: 0.5715\n",
      "2024-05-21 13:23:59 [INFO]: Epoch 004 - training loss: 22828.3269, validation loss: 0.5638\n",
      "2024-05-21 13:24:03 [INFO]: Epoch 005 - training loss: 22821.9063, validation loss: 0.5509\n",
      "2024-05-21 13:24:06 [INFO]: Epoch 006 - training loss: 22817.8396, validation loss: 0.5411\n",
      "2024-05-21 13:24:10 [INFO]: Epoch 007 - training loss: 22816.0375, validation loss: 0.5412\n",
      "2024-05-21 13:24:15 [INFO]: Epoch 008 - training loss: 22813.3457, validation loss: 0.5447\n",
      "2024-05-21 13:24:19 [INFO]: Epoch 009 - training loss: 22812.1050, validation loss: 0.5212\n",
      "2024-05-21 13:24:23 [INFO]: Epoch 010 - training loss: 22809.5779, validation loss: 0.5380\n",
      "2024-05-21 13:24:23 [INFO]: Finished training. The best model is from epoch#9.\n",
      "2024-05-21 13:24:23 [INFO]: Saved the model to tutorial_results/imputation/gp_vae/20240521_T132343/GPVAE.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "gp_vae.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The shape of gp_vae_imputation is (2398, 2, 48, 37)\n"
     ]
    }
   ],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "\n",
    "# GP-VAE has an argument to control the number of sampling times during inference\n",
    "gp_vae_results = gp_vae.predict(dataset_for_testing, n_sampling_times=2)\n",
    "gp_vae_imputation = gp_vae_results[\"imputation\"]\n",
    "\n",
    "print(f\"The shape of gp_vae_imputation is {gp_vae_imputation.shape}\")\n",
    "\n",
    "# for error calculation, we need to take the mean value of the multiple samplings for each data sample\n",
    "mean_gp_vae_imputation = gp_vae_imputation.mean(axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.4822\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    mean_gp_vae_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 An example of **BRITS** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 13:24:24 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 13:24:24 [INFO]: Model files will be saved to tutorial_results/imputation/brits/20240521_T132424\n",
      "2024-05-21 13:24:24 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/brits/20240521_T132424/tensorboard\n",
      "2024-05-21 13:24:24 [INFO]: BRITS initialized with the given hyperparameters, the number of trainable parameters: 239,344\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import BRITS\n",
    "\n",
    "# initialize the model\n",
    "brits = BRITS(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    rnn_hidden_size=128,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,\n",
    "    # set the path for saving tensorboard and trained model files \n",
    "    saving_path=\"tutorial_results/imputation/brits\", \n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 13:26:06 [INFO]: Epoch 001 - training loss: 0.9522, validation loss: 0.3980\n",
      "2024-05-21 13:27:26 [INFO]: Epoch 002 - training loss: 0.7375, validation loss: 0.3423\n",
      "2024-05-21 13:28:49 [INFO]: Epoch 003 - training loss: 0.6834, validation loss: 0.3278\n",
      "2024-05-21 13:30:15 [INFO]: Epoch 004 - training loss: 0.6578, validation loss: 0.3228\n",
      "2024-05-21 13:31:40 [INFO]: Epoch 005 - training loss: 0.6421, validation loss: 0.3166\n",
      "2024-05-21 13:33:01 [INFO]: Epoch 006 - training loss: 0.6305, validation loss: 0.3149\n",
      "2024-05-21 13:34:25 [INFO]: Epoch 007 - training loss: 0.6216, validation loss: 0.3136\n",
      "2024-05-21 13:35:48 [INFO]: Epoch 008 - training loss: 0.6142, validation loss: 0.3158\n",
      "2024-05-21 13:37:11 [INFO]: Epoch 009 - training loss: 0.6074, validation loss: 0.3169\n",
      "2024-05-21 13:38:31 [INFO]: Epoch 010 - training loss: 0.6020, validation loss: 0.3166\n",
      "2024-05-21 13:38:31 [INFO]: Exceeded the training patience. Terminating the training procedure...\n",
      "2024-05-21 13:38:31 [INFO]: Finished training. The best model is from epoch#7.\n",
      "2024-05-21 13:38:31 [INFO]: Saved the model to tutorial_results/imputation/brits/20240521_T132424/BRITS.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "brits.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "brits_results = brits.predict(dataset_for_testing)\n",
    "brits_imputation = brits_results[\"imputation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.2583\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    brits_imputation, \n",
    "    physionet2012_dataset['test_X_ori'], \n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 An example of **M-RNN** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 14:09:50 [INFO]: Using the given device: cuda:0\n",
      "2024-05-21 14:09:50 [INFO]: Model files will be saved to tutorial_results/imputation/mrnn/20240521_T140950\n",
      "2024-05-21 14:09:50 [INFO]: Tensorboard file will be saved to tutorial_results/imputation/mrnn/20240521_T140950/tensorboard\n",
      "2024-05-21 14:09:50 [INFO]: MRNN initialized with the given hyperparameters, the number of trainable parameters: 107,951\n"
     ]
    }
   ],
   "source": [
    "from pypots.optim import Adam\n",
    "from pypots.imputation import MRNN\n",
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# initialize the model\n",
    "# initialize the model\n",
    "mrnn = MRNN(\n",
    "    n_steps=physionet2012_dataset['n_steps'],\n",
    "    n_features=physionet2012_dataset['n_features'],\n",
    "    rnn_hidden_size=128,\n",
    "    batch_size=32,\n",
    "    # here we set epochs=10 for a quick demo, you can set it to 100 or more for better performance\n",
    "    epochs=10,\n",
    "    # here we set patience=3 to early stop the training if the evaluting loss doesn't decrease for 3 epoches.\n",
    "    # You can leave it to defualt as None to disable early stopping.\n",
    "    patience=3,\n",
    "    # give the optimizer. Different from torch.optim.Optimizer, you don't have to specify model's parameters when\n",
    "    # initializing pypots.optim.Optimizer. You can also leave it to default. It will initilize an Adam optimizer with lr=0.001.\n",
    "    optimizer=Adam(lr=1e-3),\n",
    "    # this num_workers argument is for torch.utils.data.Dataloader. It's the number of subprocesses to use for data loading.\n",
    "    # Leaving it to default as 0 means data loading will be in the main process, i.e. there won't be subprocesses.\n",
    "    # You can increase it to >1 if you think your dataloading is a bottleneck to your model training speed\n",
    "    num_workers=0,\n",
    "    # just leave it to default as None, PyPOTS will automatically assign the best device for you.\n",
    "    # Set it as 'cpu' if you don't have CUDA devices. You can also set it to 'cuda:0' or 'cuda:1' if you have multiple CUDA devices, even parallelly on ['cuda:0', 'cuda:1']\n",
    "    device=None,\n",
    "    # set the path for saving tensorboard and trained model files\n",
    "    saving_path=\"tutorial_results/imputation/mrnn\",\n",
    "    # only save the best model after training finished.\n",
    "    # You can also set it as \"better\" to save models performing better ever during training.\n",
    "    model_saving_strategy=\"best\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 14:10:33 [INFO]: Epoch 001 - training loss: 0.7347, validation loss: 0.9286\n",
      "2024-05-21 14:10:57 [INFO]: Epoch 002 - training loss: 0.5278, validation loss: 0.8844\n",
      "2024-05-21 14:11:18 [INFO]: Epoch 003 - training loss: 0.4864, validation loss: 0.8717\n",
      "2024-05-21 14:11:42 [INFO]: Epoch 004 - training loss: 0.4697, validation loss: 0.8623\n",
      "2024-05-21 14:12:12 [INFO]: Epoch 005 - training loss: 0.4615, validation loss: 0.8596\n",
      "2024-05-21 14:12:41 [INFO]: Epoch 006 - training loss: 0.4405, validation loss: 0.8583\n",
      "2024-05-21 14:13:01 [INFO]: Epoch 007 - training loss: 0.4361, validation loss: 0.8600\n",
      "2024-05-21 14:13:23 [INFO]: Epoch 008 - training loss: 0.4315, validation loss: 0.8633\n",
      "2024-05-21 14:13:45 [INFO]: Epoch 009 - training loss: 0.4296, validation loss: 0.8660\n",
      "2024-05-21 14:13:45 [INFO]: Exceeded the training patience. Terminating the training procedure...\n",
      "2024-05-21 14:13:45 [INFO]: Finished training. The best model is from epoch#6.\n",
      "2024-05-21 14:13:45 [INFO]: Saved the model to tutorial_results/imputation/mrnn/20240521_T140950/MRNN.pypots\n"
     ]
    }
   ],
   "source": [
    "# train the model on the training set, and validate it on the validating set to select the best model for testing in the next step\n",
    "mrnn.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "mrnn_results = mrnn.predict(dataset_for_testing)\n",
    "mrnn_imputation = mrnn_results[\"imputation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.6776\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    mrnn_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 An example of **LOCF** for imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-21 14:13:51 [INFO]: No given device, using default device: cuda\n"
     ]
    }
   ],
   "source": [
    "from pypots.imputation import LOCF\n",
    "\n",
    "# initialize the model\n",
    "locf = LOCF()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linglong/data/linglong/.conda/envs/pypots-dev/lib/python3.11/site-packages/pypots/imputation/locf/model.py:67: UserWarning: LOCF (Last Observed Carried Forward) imputation class has no parameter to train. Please run func `predict()` directly.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# LOCF doesn't need to be trained, just call the impute() function\n",
    "\n",
    "locf.fit(train_set=dataset_for_training, val_set=dataset_for_validating)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the testing stage, impute the originally-missing values and artificially-missing values in the test set\n",
    "locf_results = locf.predict(dataset_for_testing)\n",
    "locf_imputation = locf_results[\"imputation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mean absolute error: 0.4091\n"
     ]
    }
   ],
   "source": [
    "from pypots.utils.metrics import calc_mae\n",
    "\n",
    "# calculate mean absolute error on the ground truth (artificially-missing values)\n",
    "testing_mae = calc_mae(\n",
    "    locf_imputation,\n",
    "    physionet2012_dataset['test_X_ori'],\n",
    "    physionet2012_dataset['test_X_indicating_mask'],\n",
    ")\n",
    "print(f\"Testing mean absolute error: {testing_mae:.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}