tutorials/data/tutorial02_datasets.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Tutorial 02: Datasets\n",
    "\n",
    "This tutorial shows different TemporAI `Dataset`s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Skip the below cell if you are not on Google Colab / already have TemporAI installed:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install temporai\n",
    "\n",
    "# Or from the repo, for the latest version:\n",
    "# %pip install git+https://github.com/vanderschaarlab/temporai.git"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare some example data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "# Some time series data:\n",
    "time_series_df = pd.DataFrame(\n",
    "    {\n",
    "        \"sample_idx\": [\"sample_0\", \"sample_0\", \"sample_0\", \"sample_0\", \"sample_1\", \"sample_1\", \"sample_2\"],\n",
    "        \"time_idx\": [1, 2, 3, 4, 2, 4, 9],\n",
    "        \"t_feat_0\": [11, 12, 13, 14, 21, 22, 31],\n",
    "        \"t_feat_1\": [1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 3.1],\n",
    "        \"t_feat_2\": [10, 20, 30, 40, 11, 21, 111],\n",
    "    }\n",
    ")\n",
    "time_series_df.set_index(keys=[\"sample_idx\", \"time_idx\"], drop=True, inplace=True)\n",
    "\n",
    "# Some static data:\n",
    "static_df = pd.DataFrame(\n",
    "    {\n",
    "        \"s_feat_0\": [100, 200, 300],\n",
    "        \"s_feat_1\": [-1.1, -1.2, -1.3],\n",
    "        \"s_feat_2\": [0, 1, 0],\n",
    "    },\n",
    "    index=[\"sample_0\", \"sample_1\", \"sample_2\"],\n",
    ")\n",
    "\n",
    "event_df = pd.DataFrame(\n",
    "    {\n",
    "        \"e_feat_0\": [(10, True), (12, False), (13, True)],\n",
    "        \"e_feat_1\": [(10, False), (10, False), (11, True)],\n",
    "    },\n",
    "    index=[\"sample_0\", \"sample_1\", \"sample_2\"],\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preview the dataframes below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "      <th>t_feat_1</th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>1.1</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>1.2</td>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "      <td>1.3</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>1.4</td>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>2.1</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "      <td>2.2</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "      <td>3.1</td>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                     t_feat_0  t_feat_1  t_feat_2\n",
       "sample_idx time_idx                              \n",
       "sample_0   1               11       1.1        10\n",
       "           2               12       1.2        20\n",
       "           3               13       1.3        30\n",
       "           4               14       1.4        40\n",
       "sample_1   2               21       2.1        11\n",
       "           4               22       2.2        21\n",
       "sample_2   9               31       3.1       111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "time_series_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          s_feat_0  s_feat_1  s_feat_2\n",
       "sample_0       100      -1.1         0\n",
       "sample_1       200      -1.2         1\n",
       "sample_2       300      -1.3         0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "static_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>e_feat_0</th>\n",
       "      <th>e_feat_1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>(10, True)</td>\n",
       "      <td>(10, False)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>(12, False)</td>\n",
       "      <td>(10, False)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>(13, True)</td>\n",
       "      <td>(11, True)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             e_feat_0     e_feat_1\n",
       "sample_0   (10, True)  (10, False)\n",
       "sample_1  (12, False)  (10, False)\n",
       "sample_2   (13, True)   (11, True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "event_df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `CovariatesDataset`\n",
    "\n",
    "A `CovariatesDataset` contains time series and optionally static covariates only, without any predictive data\n",
    "(targets or treatments).\n",
    "\n",
    "It can be used with `preprocessing` transformations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tempor.data import dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CovariatesDataset(\n",
       "    time_series=TimeSeriesSamples([3, *, 3]),\n",
       "    static=StaticSamples([3, 3])\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize a CovariatesDataset:\n",
    "data = dataset.CovariatesDataset(\n",
    "    time_series=time_series_df,\n",
    "    static=static_df,  # Optional, can be `None`.\n",
    ")\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "      <th>t_feat_1</th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>1.1</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>1.2</td>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "      <td>1.3</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>1.4</td>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>2.1</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "      <td>2.2</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "      <td>3.1</td>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_0  t_feat_1  t_feat_2\n",
       "sample_idx time_idx                              \n",
       "sample_0   1               11       1.1        10\n",
       "           2               12       1.2        20\n",
       "           3               13       1.3        30\n",
       "           4               14       1.4        40\n",
       "sample_1   2               21       2.1        11\n",
       "           4               22       2.2        21\n",
       "sample_2   9               31       3.1       111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.time_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_0  s_feat_1  s_feat_2\n",
       "sample_idx                              \n",
       "sample_0         100      -1.1         0\n",
       "sample_1         200      -1.2         1\n",
       "sample_2         300      -1.3         0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.static"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `OneOffPredictionDataset`\n",
    "\n",
    "A `OneOffPredictionDataset` contains time series and optionally static covariates.\n",
    "\n",
    "It also needs `StaticSamples` prediction *targets* for estimators to be able to `fit` on this dataset.\n",
    "\n",
    "It can be used with `prediction.one_off` estimators. The task is to predict some one-off value for each sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OneOffPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([3, *, 3]),\n",
       "    static=StaticSamples([3, 2]),\n",
       "    predictive=OneOffPredictionTaskData(targets=StaticSamples([3, 1]))\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize a OneOffPredictionDataset:\n",
    "data = dataset.OneOffPredictionDataset(\n",
    "    time_series=time_series_df,\n",
    "    static=static_df.loc[:, :\"s_feat_1\"],  # Optional, can be `None`.\n",
    "    targets=static_df.loc[:, [\"s_feat_2\"]],  # Optional, can be `None` at inference time.\n",
    ")\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "      <th>t_feat_1</th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>1.1</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>1.2</td>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "      <td>1.3</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>1.4</td>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>2.1</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "      <td>2.2</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "      <td>3.1</td>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_0  t_feat_1  t_feat_2\n",
       "sample_idx time_idx                              \n",
       "sample_0   1               11       1.1        10\n",
       "           2               12       1.2        20\n",
       "           3               13       1.3        30\n",
       "           4               14       1.4        40\n",
       "sample_1   2               21       2.1        11\n",
       "           4               22       2.2        21\n",
       "sample_2   9               31       3.1       111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.time_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_0  s_feat_1\n",
       "sample_idx                    \n",
       "sample_0         100      -1.1\n",
       "sample_1         200      -1.2\n",
       "sample_2         300      -1.3"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.static"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_2\n",
       "sample_idx          \n",
       "sample_0           0\n",
       "sample_1           1\n",
       "sample_2           0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.targets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `TemporalPredictionDataset`\n",
    "\n",
    "A `TemporalPredictionDataset` contains time series and optionally static covariates.\n",
    "\n",
    "It also needs `TimeSeriesSamples` prediction *targets* for estimators to be able to `fit` on this dataset.\n",
    "\n",
    "It can be used with `prediction.temporal` estimators. The task is to predict some time series for each sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TemporalPredictionDataset(\n",
       "    time_series=TimeSeriesSamples([3, *, 2]),\n",
       "    static=StaticSamples([3, 3]),\n",
       "    predictive=TemporalPredictionTaskData(\n",
       "        targets=TimeSeriesSamples([3, *, 1])\n",
       "    )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize a TemporalPredictionDataset:\n",
    "data = dataset.TemporalPredictionDataset(\n",
    "    time_series=time_series_df.loc[:, :\"t_feat_1\"],\n",
    "    static=static_df,  # Optional, can be `None`.\n",
    "    targets=time_series_df.loc[:, [\"t_feat_2\"]],  # Optional, can be `None` at inference time.\n",
    ")\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "      <th>t_feat_1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>1.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>1.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "      <td>1.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>1.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>2.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "      <td>2.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "      <td>3.1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_0  t_feat_1\n",
       "sample_idx time_idx                    \n",
       "sample_0   1               11       1.1\n",
       "           2               12       1.2\n",
       "           3               13       1.3\n",
       "           4               14       1.4\n",
       "sample_1   2               21       2.1\n",
       "           4               22       2.2\n",
       "sample_2   9               31       3.1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.time_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_0  s_feat_1  s_feat_2\n",
       "sample_idx                              \n",
       "sample_0         100      -1.1         0\n",
       "sample_1         200      -1.2         1\n",
       "sample_2         300      -1.3         0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.static"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_2\n",
       "sample_idx time_idx          \n",
       "sample_0   1               10\n",
       "           2               20\n",
       "           3               30\n",
       "           4               40\n",
       "sample_1   2               11\n",
       "           4               21\n",
       "sample_2   9              111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.targets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `TimeToEventAnalysisDataset`\n",
    "\n",
    "A `TimeToEventAnalysisDataset` contains time series and optionally static covariates.\n",
    "\n",
    "It also needs `EventSamples` prediction *targets* for estimators to be able to `fit` on this dataset.\n",
    "\n",
    "It can be used with `time_to_event` estimators. The task is to predict risk scores for each sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TimeToEventAnalysisDataset(\n",
       "    time_series=TimeSeriesSamples([3, *, 3]),\n",
       "    static=StaticSamples([3, 3]),\n",
       "    predictive=TimeToEventAnalysisTaskData(targets=EventSamples([3, 2]))\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize a TimeToEventAnalysisDataset:\n",
    "data = dataset.TimeToEventAnalysisDataset(\n",
    "    time_series=time_series_df,\n",
    "    static=static_df,  # Optional, can be `None`.\n",
    "    targets=event_df,  # Optional, can be `None` at inference time.\n",
    ")\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "      <th>t_feat_1</th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>1.1</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>1.2</td>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "      <td>1.3</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>1.4</td>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>2.1</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "      <td>2.2</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "      <td>3.1</td>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_0  t_feat_1  t_feat_2\n",
       "sample_idx time_idx                              \n",
       "sample_0   1               11       1.1        10\n",
       "           2               12       1.2        20\n",
       "           3               13       1.3        30\n",
       "           4               14       1.4        40\n",
       "sample_1   2               21       2.1        11\n",
       "           4               22       2.2        21\n",
       "sample_2   9               31       3.1       111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.time_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_0  s_feat_1  s_feat_2\n",
       "sample_idx                              \n",
       "sample_0         100      -1.1         0\n",
       "sample_1         200      -1.2         1\n",
       "sample_2         300      -1.3         0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.static"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">EventSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>e_feat_0</th>\n",
       "      <th>e_feat_1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>(10, True)</td>\n",
       "      <td>(10, False)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>(12, False)</td>\n",
       "      <td>(10, False)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>(13, True)</td>\n",
       "      <td>(11, True)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "EventSamples with data:\n",
       "               e_feat_0     e_feat_1\n",
       "sample_idx                          \n",
       "sample_0     (10, True)  (10, False)\n",
       "sample_1    (12, False)  (10, False)\n",
       "sample_2     (13, True)   (11, True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.targets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `OneOffTreatmentEffectsDataset`\n",
    "\n",
    "A `OneOffTreatmentEffectsDataset` contains time series and optionally static covariates.\n",
    "\n",
    "It also needs `TimeSeriesSamples` prediction *targets* and `EventSamples` treatments\n",
    "for estimators to be able to `fit` on this dataset.\n",
    "\n",
    "It can be used with `treatments.one_off` estimators.\n",
    "The task is to predict a time series counterfactual outcome based on a one-off treatment event."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OneOffTreatmentEffectsDataset(\n",
       "    time_series=TimeSeriesSamples([3, *, 2]),\n",
       "    static=StaticSamples([3, 3]),\n",
       "    predictive=OneOffTreatmentEffectsTaskData(\n",
       "        targets=TimeSeriesSamples([3, *, 1]),\n",
       "        treatments=EventSamples([3, 1])\n",
       "    )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize a TimeToEventAnalysisDataset:\n",
    "data = dataset.OneOffTreatmentEffectsDataset(\n",
    "    time_series=time_series_df.loc[:, :\"t_feat_1\"],\n",
    "    static=static_df,  # Optional, can be `None`.\n",
    "    targets=time_series_df.loc[:, [\"t_feat_2\"]],  # Optional, can be `None` at inference time.\n",
    "    treatments=event_df.loc[:, [\"e_feat_0\"]],\n",
    ")\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "      <th>t_feat_1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>1.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>1.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "      <td>1.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>1.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>2.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "      <td>2.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "      <td>3.1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_0  t_feat_1\n",
       "sample_idx time_idx                    \n",
       "sample_0   1               11       1.1\n",
       "           2               12       1.2\n",
       "           3               13       1.3\n",
       "           4               14       1.4\n",
       "sample_1   2               21       2.1\n",
       "           4               22       2.2\n",
       "sample_2   9               31       3.1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.time_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_0  s_feat_1  s_feat_2\n",
       "sample_idx                              \n",
       "sample_0         100      -1.1         0\n",
       "sample_1         200      -1.2         1\n",
       "sample_2         300      -1.3         0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.static"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_2\n",
       "sample_idx time_idx          \n",
       "sample_0   1               10\n",
       "           2               20\n",
       "           3               30\n",
       "           4               40\n",
       "sample_1   2               11\n",
       "           4               21\n",
       "sample_2   9              111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">EventSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>e_feat_0</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>(10, True)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>(12, False)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>(13, True)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "EventSamples with data:\n",
       "               e_feat_0\n",
       "sample_idx             \n",
       "sample_0     (10, True)\n",
       "sample_1    (12, False)\n",
       "sample_2     (13, True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.treatments"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `TemporalTreatmentEffectsDataset`\n",
    "\n",
    "A `TemporalTreatmentEffectsDataset` contains time series and optionally static covariates.\n",
    "\n",
    "It also needs `TimeSeriesSamples` prediction *targets* and `TimeSeriesSamples` treatments\n",
    "for estimators to be able to `fit` on this dataset.\n",
    "\n",
    "It can be used with `treatments.temporal` estimators.\n",
    "The task is to predict a time series counterfactual outcome based on a time series treatment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TemporalTreatmentEffectsDataset(\n",
       "    time_series=TimeSeriesSamples([3, *, 1]),\n",
       "    static=StaticSamples([3, 3]),\n",
       "    predictive=TemporalTreatmentEffectsTaskData(\n",
       "        targets=TimeSeriesSamples([3, *, 1]),\n",
       "        treatments=TimeSeriesSamples([3, *, 1])\n",
       "    )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize a TimeToEventAnalysisDataset:\n",
    "data = dataset.TemporalTreatmentEffectsDataset(\n",
    "    time_series=time_series_df.loc[:, :\"t_feat_0\"],\n",
    "    static=static_df,  # Optional, can be `None`.\n",
    "    targets=time_series_df.loc[:, [\"t_feat_1\"]],  # Optional, can be `None` at inference time.\n",
    "    treatments=time_series_df.loc[:, [\"t_feat_2\"]],\n",
    ")\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_0</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>31</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_0\n",
       "sample_idx time_idx          \n",
       "sample_0   1               11\n",
       "           2               12\n",
       "           3               13\n",
       "           4               14\n",
       "sample_1   2               21\n",
       "           4               22\n",
       "sample_2   9               31"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.time_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">StaticSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s_feat_0</th>\n",
       "      <th>s_feat_1</th>\n",
       "      <th>s_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>sample_0</th>\n",
       "      <td>100</td>\n",
       "      <td>-1.1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_1</th>\n",
       "      <td>200</td>\n",
       "      <td>-1.2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <td>300</td>\n",
       "      <td>-1.3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "StaticSamples with data:\n",
       "            s_feat_0  s_feat_1  s_feat_2\n",
       "sample_idx                              \n",
       "sample_0         100      -1.1         0\n",
       "sample_1         200      -1.2         1\n",
       "sample_2         300      -1.3         0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.static"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>1.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>2.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>3.1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_1\n",
       "sample_idx time_idx          \n",
       "sample_0   1              1.1\n",
       "           2              1.2\n",
       "           3              1.3\n",
       "           4              1.4\n",
       "sample_1   2              2.1\n",
       "           4              2.2\n",
       "sample_2   9              3.1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p><span style=\"font-family: monospace;\">TimeSeriesSamples</span> with data:</p><div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>t_feat_2</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_idx</th>\n",
       "      <th>time_idx</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">sample_0</th>\n",
       "      <th>1</th>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">sample_1</th>\n",
       "      <th>2</th>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_2</th>\n",
       "      <th>9</th>\n",
       "      <td>111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "TimeSeriesSamples with data:\n",
       "                     t_feat_2\n",
       "sample_idx time_idx          \n",
       "sample_0   1               10\n",
       "           2               20\n",
       "           3               30\n",
       "           4               40\n",
       "sample_1   2               11\n",
       "           4               21\n",
       "sample_2   9              111"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.predictive.treatments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🎉 Congratulations!\n",
    "\n",
    "Congratulations on completing this notebook tutorial! If you enjoyed this and would like to join the movement towards *Machine learning and AI for Medicine*, you can do so in the following ways!\n",
    "\n",
    "\n",
    "\n",
    "### ⭐ Star [TemporAI](https://github.com/vanderschaarlab/temporai) on GitHub\n",
    "\n",
    "- The easiest way to help our community is by just starring the repos! This helps raise awareness of the tools we're building.\n",
    "\n",
    "\n",
    "\n",
    "### Check out other projects from [vanderschaarlab](https://github.com/vanderschaarlab)\n",
    "- 📝 [HyperImpute](https://github.com/vanderschaarlab/hyperimpute)\n",
    "- 📊 [AutoPrognosis](https://github.com/vanderschaarlab/autoprognosis)\n",
    "- 🤖 [SynthCity](https://github.com/vanderschaarlab/synthcity)\n",
    " "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "RE_py38_temporai",
   "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}