{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp transforms.splitting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Splitting\n",
    "> Data Splitting Transforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import numpy as np\n",
    "import math\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from collections import defaultdict\n",
    "import functools\n",
    "import operator\n",
    "from typing import Any, Iterable, Optional, Tuple, Union\n",
    "\n",
    "from scipy.sparse import coo_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def _subset_interactions(mat, idxs):\n",
    "    idxs = np.array(idxs)\n",
    "\n",
    "    coo_mat = coo_matrix(\n",
    "        (mat.data[idxs], (mat.row[idxs], mat.col[idxs])),\n",
    "        shape=(mat.shape[0], mat.shape[1])\n",
    "    )\n",
    "\n",
    "    return coo_mat\n",
    "\n",
    "\n",
    "def _validate_val_p_and_test_p(val_p: float, test_p: float) -> None:\n",
    "    validate_and_test_p = val_p + test_p\n",
    "\n",
    "    if val_p >= 1 or val_p < 0:\n",
    "        raise ValueError('``val_p`` must be in the range [0, 1).')\n",
    "    if test_p >= 1 or test_p < 0:\n",
    "        raise ValueError('``test_p`` must be in the range [0, 1).')\n",
    "    if validate_and_test_p >= 1 or validate_and_test_p <= 0:\n",
    "        raise ValueError('The sum of ``val_p`` and ``test_p`` must be in the range (0, 1).')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def random_split(mat, val_p = 0.0, test_p = 0.2, seed = 42):\n",
    "    \"\"\"Randomly split interactions into training, validation, and testing sets.\"\"\"\n",
    "\n",
    "    _validate_val_p_and_test_p(val_p=val_p, test_p=test_p)\n",
    "\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    num_interactions = mat.nnz\n",
    "\n",
    "    shuffle_indices = np.arange(num_interactions)\n",
    "    np.random.shuffle(shuffle_indices)\n",
    "\n",
    "    interactions = _subset_interactions(mat=mat, idxs=shuffle_indices)\n",
    "\n",
    "    validate_and_test_p = val_p + test_p\n",
    "    validate_cutoff = int((1.0 - validate_and_test_p) * num_interactions)\n",
    "    test_cutoff = int((1.0 - test_p) * num_interactions)\n",
    "\n",
    "    train_idxs = np.arange(validate_cutoff)\n",
    "    validate_idxs = np.arange(validate_cutoff, test_cutoff)\n",
    "    test_idxs = np.arange(test_cutoff, num_interactions)\n",
    "\n",
    "    train_interactions = _subset_interactions(mat=mat, idxs=train_idxs)\n",
    "    test_interactions = _subset_interactions(mat=mat, idxs=test_idxs)\n",
    "\n",
    "    if val_p > 0:\n",
    "        validate_interactions = _subset_interactions(mat=mat, idxs=validate_idxs)\n",
    "\n",
    "        return train_interactions, validate_interactions, test_interactions\n",
    "    else:\n",
    "        return train_interactions, test_interactions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "def test_random_split():\n",
    "\n",
    "    interactions_to_split_df = pd.DataFrame(\n",
    "        data={\n",
    "        'user_id': [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4, 4],\n",
    "        'item_id': [0, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 4, 1, 2, 4, 5],\n",
    "        'rating': [1, 2, 3, 4, 5, 4, 3, 2, 1, 1, 2, 3, 4, 2, 3, 4, 5, 1, 5, 4, 2, 3, 5, 4]\n",
    "        }\n",
    "    )\n",
    "    interactions_to_split=coo_matrix(\n",
    "        (\n",
    "            interactions_to_split_df['rating'],\n",
    "            (interactions_to_split_df['user_id'], interactions_to_split_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split_df.user_id.nunique(), interactions_to_split_df.item_id.nunique()),\n",
    "    )\n",
    "\n",
    "    train_expected_df = pd.DataFrame(\n",
    "        data={\n",
    "            'user_id': [0, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 0, 0, 0, 0],\n",
    "            'item_id': [0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 5, 6, 7, 8],\n",
    "            'rating': [1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5, 4, 4, 3, 2, 1], \n",
    "        }\n",
    "    )\n",
    "    train_expected=coo_matrix(\n",
    "        (\n",
    "            train_expected_df['rating'],\n",
    "            (train_expected_df['user_id'], train_expected_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split.shape[0], interactions_to_split.shape[1]),\n",
    "    )\n",
    "\n",
    "    validate_expected_df = pd.DataFrame(\n",
    "        data={'user_id': [2, 3, 3], 'item_id': [4, 1, 2], 'rating': [5, 1, 5]}\n",
    "    )\n",
    "    validate_expected=coo_matrix(\n",
    "        (\n",
    "            validate_expected_df['rating'],\n",
    "            (validate_expected_df['user_id'], validate_expected_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split.shape[0], interactions_to_split.shape[1]),\n",
    "    )\n",
    "\n",
    "    test_expected_df = pd.DataFrame(\n",
    "        data={\n",
    "            'user_id': [3, 4, 4, 4, 4],\n",
    "            'item_id': [4, 1, 2, 4, 5],\n",
    "            'rating': [4, 2, 3, 5, 4],\n",
    "        }\n",
    "    )\n",
    "    test_expected=coo_matrix(\n",
    "        (\n",
    "            test_expected_df['rating'],\n",
    "            (test_expected_df['user_id'], test_expected_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split.shape[0], interactions_to_split.shape[1]),\n",
    "    )\n",
    "\n",
    "    (train_actual, validate_actual, test_actual) = random_split(\n",
    "        mat=interactions_to_split, val_p=0.1, test_p=0.2, seed=42\n",
    "    )\n",
    "\n",
    "    np.testing.assert_array_equal(train_actual.toarray(), train_expected.toarray())\n",
    "    np.testing.assert_array_equal(\n",
    "        validate_actual.toarray(), validate_expected.toarray()\n",
    "    )\n",
    "    np.testing.assert_array_equal(test_actual.toarray(), test_expected.toarray())\n",
    "\n",
    "    assert (\n",
    "        train_actual.shape[0]\n",
    "        == train_expected.shape[0]\n",
    "        == validate_actual.shape[0]\n",
    "        == validate_expected.shape[0]\n",
    "        == test_actual.shape[0]\n",
    "        == test_expected.shape[0]\n",
    "    )\n",
    "\n",
    "    assert (\n",
    "        train_actual.shape[1]\n",
    "        == train_expected.shape[1]\n",
    "        == validate_actual.shape[1]\n",
    "        == validate_expected.shape[1]\n",
    "        == test_actual.shape[1]\n",
    "        == test_expected.shape[1]\n",
    "    )\n",
    "\n",
    "    assert (\n",
    "        type(train_actual)\n",
    "        == type(train_expected)\n",
    "        == type(validate_actual)\n",
    "        == type(validate_expected)\n",
    "        == type(test_actual)\n",
    "        == type(test_expected)\n",
    "    )\n",
    "\n",
    "test_random_split()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stratified split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def _stratified_split(mat, test_p, seed):\n",
    "    users = mat.row\n",
    "    unique_users = set(users)\n",
    "\n",
    "    # while we should be able to run ``np.where(users == user)[0]`` to find all items each user\n",
    "    # interacted with, by building up a dictionary to get these values instead, we can achieve the\n",
    "    # same result in O(N) complexity rather than O(M * N), a nice timesave to have when working with\n",
    "    # larger datasets\n",
    "    all_idxs_for_users_dict = defaultdict(list)\n",
    "    for idx, user in enumerate(users):\n",
    "        all_idxs_for_users_dict[user].append(idx)\n",
    "\n",
    "    test_idxs = [\n",
    "        _stratified_split_parallel_worker(idxs_to_split=all_idxs_for_users_dict[user],\n",
    "                                            test_p=test_p,\n",
    "                                            seed=(seed + user))\n",
    "        for user in unique_users\n",
    "    ]\n",
    "\n",
    "    # reduce the list of lists down to a 1-d list\n",
    "    test_idxs = functools.reduce(operator.iconcat, test_idxs, [])\n",
    "    # find all indices not in test set - they are now train\n",
    "    train_idxs = list(set(range(len(users))) - set(test_idxs))\n",
    "\n",
    "    train_interactions = _subset_interactions(mat=mat,\n",
    "                                              idxs=train_idxs)\n",
    "    test_interactions = _subset_interactions(mat=mat,\n",
    "                                             idxs=test_idxs)\n",
    "\n",
    "    return train_interactions, test_interactions\n",
    "\n",
    "\n",
    "def _stratified_split_parallel_worker(idxs_to_split: Iterable[Any],\n",
    "                                      test_p: float, seed: int) -> np.array:\n",
    "    _, test_idxs = train_test_split(idxs_to_split,\n",
    "                                    test_size=test_p,\n",
    "                                    random_state=seed,\n",
    "                                    shuffle=True,\n",
    "                                    stratify=np.ones_like(idxs_to_split))\n",
    "\n",
    "    return test_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def stratified_split(mat, val_p = 0.0, test_p = 0.2, seed = 42):\n",
    "    \"\"\"Split into train, validate, and test datasets in a stratified\n",
    "    manner such that each user appears at least once in each of the datasets.\n",
    "    This split guarantees that every user will be represented in the training, validation, and\n",
    "    testing datasets given they appear at least three times. If ``val_p ==\n",
    "    0``, they will appear in the training and testing datasets given they appear at least two times.\n",
    "    If a user appears fewer than this number of times, a ``ValueError`` will\n",
    "    be raised.\"\"\"\n",
    "\n",
    "    _validate_val_p_and_test_p(val_p=val_p, test_p=test_p)\n",
    "\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    train, test = _stratified_split(mat=mat,\n",
    "                                    test_p=test_p,\n",
    "                                    seed=seed)\n",
    "\n",
    "    if val_p > 0:\n",
    "        train, validate = _stratified_split(mat=train,\n",
    "                                            test_p=val_p / (1 - test_p),\n",
    "                                            seed=seed)\n",
    "\n",
    "        return train, validate, test\n",
    "    else:\n",
    "        return train, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "\n",
    "def test_stratified_split():\n",
    "\n",
    "    interactions_to_split_df = pd.DataFrame(\n",
    "        data={\n",
    "        'user_id': [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4, 4],\n",
    "        'item_id': [0, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 4, 1, 2, 4, 5],\n",
    "        'rating': [1, 2, 3, 4, 5, 4, 3, 2, 1, 1, 2, 3, 4, 2, 3, 4, 5, 1, 5, 4, 2, 3, 5, 4]\n",
    "        }\n",
    "    )\n",
    "    interactions_to_split=coo_matrix(\n",
    "        (\n",
    "            interactions_to_split_df['rating'],\n",
    "            (interactions_to_split_df['user_id'], interactions_to_split_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split_df.user_id.nunique(), interactions_to_split_df.item_id.nunique()),\n",
    "    )\n",
    "\n",
    "    train_expected_df = pd.DataFrame(\n",
    "        data={\n",
    "            'user_id': [0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 3, 4, 4],\n",
    "            'item_id': [0, 1, 2, 4, 5, 7, 3, 4, 3, 4, 1, 1, 5],\n",
    "            'rating': [1, 2, 3, 5, 4, 2, 3, 4, 4, 5, 1, 2, 4],  \n",
    "        }\n",
    "    )\n",
    "    train_expected=coo_matrix(\n",
    "        (\n",
    "            train_expected_df['rating'],\n",
    "            (train_expected_df['user_id'], train_expected_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split.shape[0], interactions_to_split.shape[1]),\n",
    "    )\n",
    "\n",
    "    validate_expected_df = pd.DataFrame(\n",
    "        data={'user_id': [0, 1, 2, 3, 4], 'item_id': [8, 2, 2, 2, 4], 'rating': [1, 2, 3, 5, 5]}\n",
    "    )\n",
    "    validate_expected=coo_matrix(\n",
    "        (\n",
    "            validate_expected_df['rating'],\n",
    "            (validate_expected_df['user_id'], validate_expected_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split.shape[0], interactions_to_split.shape[1]),\n",
    "    )\n",
    "\n",
    "    test_expected_df = pd.DataFrame(\n",
    "        data={\n",
    "            'user_id': [0, 0, 1, 2, 3, 4],\n",
    "            'item_id': [3, 6, 1, 1, 4, 2],\n",
    "            'rating': [4, 3, 1, 2, 4, 3],\n",
    "        }\n",
    "    )\n",
    "    test_expected=coo_matrix(\n",
    "        (\n",
    "            test_expected_df['rating'],\n",
    "            (test_expected_df['user_id'], test_expected_df['item_id']),\n",
    "        ),\n",
    "        shape=(interactions_to_split.shape[0], interactions_to_split.shape[1]),\n",
    "    )\n",
    "\n",
    "    (train_actual, validate_actual, test_actual) = stratified_split(\n",
    "        mat=interactions_to_split, val_p=0.1, test_p=0.2, seed=42\n",
    "    )\n",
    "\n",
    "    np.testing.assert_array_equal(train_actual.toarray(), train_expected.toarray())\n",
    "    np.testing.assert_array_equal(\n",
    "        validate_actual.toarray(), validate_expected.toarray()\n",
    "    )\n",
    "    np.testing.assert_array_equal(test_actual.toarray(), test_expected.toarray())\n",
    "\n",
    "    assert (\n",
    "        train_actual.shape[0]\n",
    "        == train_expected.shape[0]\n",
    "        == validate_actual.shape[0]\n",
    "        == validate_expected.shape[0]\n",
    "        == test_actual.shape[0]\n",
    "        == test_expected.shape[0]\n",
    "    )\n",
    "\n",
    "    assert (\n",
    "        train_actual.shape[1]\n",
    "        == train_expected.shape[1]\n",
    "        == validate_actual.shape[1]\n",
    "        == validate_expected.shape[1]\n",
    "        == test_actual.shape[1]\n",
    "        == test_expected.shape[1]\n",
    "    )\n",
    "\n",
    "    assert (\n",
    "        type(train_actual)\n",
    "        == type(train_expected)\n",
    "        == type(validate_actual)\n",
    "        == type(validate_expected)\n",
    "        == type(test_actual)\n",
    "        == type(test_expected)\n",
    "    )\n",
    "\n",
    "\n",
    "test_stratified_split()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split by Ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def groupby_user(user_indices):\n",
    "    users, user_position, user_counts = np.unique(user_indices,\n",
    "                                                  return_inverse=True,\n",
    "                                                  return_counts=True)\n",
    "    user_split_indices = np.split(np.argsort(user_position, kind=\"mergesort\"),\n",
    "                                  np.cumsum(user_counts)[:-1])\n",
    "    return user_split_indices\n",
    "\n",
    "\n",
    "def _pad_unknown_item(data_list):\n",
    "    train_data, test_data = data_list\n",
    "    n_items = train_data.item.nunique()\n",
    "    unique_items = set(train_data.item.tolist())\n",
    "    test_data.loc[~test_data.item.isin(unique_items), \"item\"] = n_items\n",
    "    return train_data, test_data\n",
    "\n",
    "\n",
    "def _filter_unknown_user_item(data_list):\n",
    "    train_data, test_data = data_list\n",
    "    unique_values = dict(user=set(train_data.user.tolist()),\n",
    "                         item=set(train_data.item.tolist()))\n",
    "\n",
    "    print(f\"test data size before filtering: {len(test_data)}\")\n",
    "    out_of_bounds_row_indices = set()\n",
    "    for col in [\"user\", \"item\"]:\n",
    "        for j, val in enumerate(test_data[col]):\n",
    "            if val not in unique_values[col]:\n",
    "                out_of_bounds_row_indices.add(j)\n",
    "\n",
    "    mask = np.arange(len(test_data))\n",
    "    test_data_clean = test_data[~np.isin(mask, list(out_of_bounds_row_indices))]\n",
    "    print(f\"test data size after filtering: {len(test_data_clean)}\")\n",
    "    return train_data, test_data_clean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def split_by_ratio(data, shuffle=False, test_size=None, pad_unknown=True,\n",
    "                   filter_unknown=False, seed=42):\n",
    "    np.random.seed(seed)\n",
    "    data = data.copy()\n",
    "    n_users = data.user.nunique()\n",
    "    user_indices = data.user.to_numpy()\n",
    "    user_split_indices = groupby_user(user_indices)\n",
    "\n",
    "    split_indices_all = [[], []]\n",
    "    for u in range(n_users):\n",
    "        u_data = user_split_indices[u]\n",
    "        u_data_len = len(u_data)\n",
    "        if u_data_len <= 3:   # keep items of rare users in trainset\n",
    "            split_indices_all[0].extend(u_data)\n",
    "        else:\n",
    "            train_threshold = round((1 - test_size) * u_data_len)\n",
    "            split_indices_all[0].extend(list(u_data[:train_threshold]))\n",
    "            split_indices_all[1].extend(list(u_data[train_threshold:]))\n",
    "\n",
    "    if shuffle:\n",
    "        split_data_all = tuple(\n",
    "            np.random.permutation(data[idx]) for idx in split_indices_all\n",
    "        )\n",
    "    else:\n",
    "        split_data_all = list(data.iloc[idx] for idx in split_indices_all)\n",
    "\n",
    "    if pad_unknown:\n",
    "        split_data_all = _pad_unknown_item(split_data_all)\n",
    "    elif filter_unknown:\n",
    "        split_data_all = _filter_unknown_user_item(split_data_all)\n",
    "    return split_data_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-2ba60b87-c2a3-4c5b-a11b-41d09be90b5a\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <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>user</th>\n",
       "      <th>item</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-2ba60b87-c2a3-4c5b-a11b-41d09be90b5a')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-2ba60b87-c2a3-4c5b-a11b-41d09be90b5a button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-2ba60b87-c2a3-4c5b-a11b-41d09be90b5a');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "   user  item\n",
       "0     1     1\n",
       "1     1     2\n",
       "2     1     3\n",
       "3     1     2\n",
       "4     1     2\n",
       "5     2     1\n",
       "6     2     2\n",
       "7     3     3"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "df = pd.DataFrame.from_dict(\n",
    "    {\n",
    "        'user':[1,1,1,1,1,2,2,3],\n",
    "        'item':[1,2,3,2,2,1,2,3]\n",
    "    }\n",
    ")\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train:\n",
      "   user  item\n",
      "0     1     1\n",
      "1     1     2\n",
      "2     1     3\n",
      "3     1     2\n",
      "5     2     1\n",
      "6     2     2\n",
      "7     3     3\n",
      "\n",
      "test:\n",
      "   user  item\n",
      "4     1     2\n"
     ]
    }
   ],
   "source": [
    "train, test = split_by_ratio(df, shuffle=False, test_size=0.2, pad_unknown=True, filter_unknown=False)\n",
    "print(\"train:\\n{}\\n\\ntest:\\n{}\".format(train,test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train:\n",
      "   user  item\n",
      "0     1     1\n",
      "1     1     2\n",
      "2     1     3\n",
      "5     2     1\n",
      "6     2     2\n",
      "7     3     3\n",
      "\n",
      "test:\n",
      "   user  item\n",
      "3     1     2\n",
      "4     1     2\n"
     ]
    }
   ],
   "source": [
    "train, test = split_by_ratio(df, shuffle=False, test_size=0.4, pad_unknown=True, filter_unknown=True)\n",
    "print(\"train:\\n{}\\n\\ntest:\\n{}\".format(train,test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Last-session-out Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def clean_split(train, test):\n",
    "    \"\"\"\n",
    "    Remove new items from the test set.\n",
    "    :param train: The training set.\n",
    "    :param test: The test set.\n",
    "    :return: The cleaned training and test sets.\n",
    "    \"\"\"\n",
    "    train_items = set()\n",
    "    train['sequence'].apply(lambda seq: train_items.update(set(seq)))\n",
    "    test['sequence'] = test['sequence'].apply(lambda seq: [it for it in seq if it in train_items])\n",
    "    return train, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def last_session_out_split(data,\n",
    "                           user_key='user_id',\n",
    "                           session_key='session_id',\n",
    "                           time_key='ts'):\n",
    "    \"\"\"\n",
    "    Assign the last session of every user to the test set and the remaining ones to the training set\n",
    "    \"\"\"\n",
    "    sessions = data.sort_values(by=[user_key, time_key]).groupby(user_key)[session_key]\n",
    "    last_session = sessions.last()\n",
    "    train = data[~data.session_id.isin(last_session.values)].copy()\n",
    "    test = data[data.session_id.isin(last_session.values)].copy()\n",
    "    train, test = clean_split(train, test)\n",
    "    return train, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-b26c0c86-6ce6-4ae4-b658-476e7b036ebd\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <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>session_id</th>\n",
       "      <th>sequence</th>\n",
       "      <th>ts</th>\n",
       "      <th>user_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>357</td>\n",
       "      <td>[793, 3489]</td>\n",
       "      <td>1421003874</td>\n",
       "      <td>4296</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>359</td>\n",
       "      <td>[1762]</td>\n",
       "      <td>1421018535</td>\n",
       "      <td>4296</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>394</td>\n",
       "      <td>[1256]</td>\n",
       "      <td>1421007470</td>\n",
       "      <td>30980</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4127</td>\n",
       "      <td>[1948, 1364, 2060, 1115, 6488, 2060]</td>\n",
       "      <td>1421416896</td>\n",
       "      <td>28117</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>6400</td>\n",
       "      <td>[687, 1394]</td>\n",
       "      <td>1420807778</td>\n",
       "      <td>35247</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-b26c0c86-6ce6-4ae4-b658-476e7b036ebd')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-b26c0c86-6ce6-4ae4-b658-476e7b036ebd button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-b26c0c86-6ce6-4ae4-b658-476e7b036ebd');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "   session_id                              sequence          ts  user_id\n",
       "0         357                           [793, 3489]  1421003874     4296\n",
       "1         359                                [1762]  1421018535     4296\n",
       "2         394                                [1256]  1421007470    30980\n",
       "3        4127  [1948, 1364, 2060, 1115, 6488, 2060]  1421416896    28117\n",
       "4        6400                           [687, 1394]  1420807778    35247"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame.from_dict({\n",
    "    'session_id': [357,359,394,4127,6400],\n",
    "    'sequence': [[793, 3489],[1762],[1256],\n",
    "                 [1948, 1364, 2060, 1115, 6488, 2060],\n",
    "                 [687, 1394]],\n",
    "    'ts': [1421003874, 1421018535, 1421007470,\n",
    "           1421416896, 1420807778],\n",
    "    'user_id': [4296, 4296, 30980, 28117, 35247]\n",
    "})\n",
    "\n",
    "df\t\t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-bd45d417-2f71-4392-9af7-1b70527ca3ff\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <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>session_id</th>\n",
       "      <th>sequence</th>\n",
       "      <th>ts</th>\n",
       "      <th>user_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>357</td>\n",
       "      <td>[793, 3489]</td>\n",
       "      <td>1421003874</td>\n",
       "      <td>4296</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-bd45d417-2f71-4392-9af7-1b70527ca3ff')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-bd45d417-2f71-4392-9af7-1b70527ca3ff button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-bd45d417-2f71-4392-9af7-1b70527ca3ff');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "   session_id     sequence          ts  user_id\n",
       "0         357  [793, 3489]  1421003874     4296"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data, test_data = last_session_out_split(df)\n",
    "train_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def process_split_ratio(ratio):\n",
    "    \"\"\"Generate split ratio lists.\n",
    "    Args:\n",
    "        ratio (float or list): a float number that indicates split ratio or a list of float\n",
    "        numbers that indicate split ratios (if it is a multi-split).\n",
    "    Returns:\n",
    "        tuple:\n",
    "        - bool: A boolean variable multi that indicates if the splitting is multi or single.\n",
    "        - list: A list of normalized split ratios.\n",
    "    \"\"\"\n",
    "    if isinstance(ratio, float):\n",
    "        if ratio <= 0 or ratio >= 1:\n",
    "            raise ValueError(\"Split ratio has to be between 0 and 1\")\n",
    "\n",
    "        multi = False\n",
    "    elif isinstance(ratio, list):\n",
    "        if any([x <= 0 for x in ratio]):\n",
    "            raise ValueError(\n",
    "                \"All split ratios in the ratio list should be larger than 0.\"\n",
    "            )\n",
    "\n",
    "        # normalize split ratios if they are not summed to 1\n",
    "        if math.fsum(ratio) != 1.0:\n",
    "            ratio = [x / math.fsum(ratio) for x in ratio]\n",
    "\n",
    "        multi = True\n",
    "    else:\n",
    "        raise TypeError(\"Split ratio should be either float or a list of floats.\")\n",
    "\n",
    "    return multi, ratio\n",
    "\n",
    "\n",
    "def _get_column_name(name, col_user, col_item):\n",
    "    if name == \"user\":\n",
    "        return col_user\n",
    "    elif name == \"item\":\n",
    "        return col_item\n",
    "    else:\n",
    "        raise ValueError(\"name should be either 'user' or 'item'.\")\n",
    "\n",
    "\n",
    "def split_pandas_data_with_ratios(data, ratios, seed=42, shuffle=False):\n",
    "    \"\"\"Helper function to split pandas DataFrame with given ratios\n",
    "    .. note::\n",
    "        Implementation referenced from `this source <https://stackoverflow.com/questions/38250710/how-to-split-data-into-3-sets-train-validation-and-test>`_.\n",
    "    Args:\n",
    "        data (pandas.DataFrame): Pandas data frame to be split.\n",
    "        ratios (list of floats): list of ratios for split. The ratios have to sum to 1.\n",
    "        seed (int): random seed.\n",
    "        shuffle (bool): whether data will be shuffled when being split.\n",
    "    Returns:\n",
    "        list: List of pd.DataFrame split by the given specifications.\n",
    "    \"\"\"\n",
    "    if math.fsum(ratios) != 1.0:\n",
    "        raise ValueError(\"The ratios have to sum to 1\")\n",
    "\n",
    "    split_index = np.cumsum(ratios).tolist()[:-1]\n",
    "\n",
    "    if shuffle:\n",
    "        data = data.sample(frac=1, random_state=seed)\n",
    "\n",
    "    splits = np.split(data, [round(x * len(data)) for x in split_index])\n",
    "\n",
    "    # Add split index (this makes splitting by group more efficient).\n",
    "    for i in range(len(ratios)):\n",
    "        splits[i][\"split_index\"] = i\n",
    "\n",
    "    return splits\n",
    "\n",
    "\n",
    "def min_rating_filter_pandas(\n",
    "    data,\n",
    "    min_rating=1,\n",
    "    filter_by=\"user\",\n",
    "    col_user='USERID',\n",
    "    col_item='ITEMID',\n",
    "):\n",
    "    \"\"\"Filter rating DataFrame for each user with minimum rating.\n",
    "    Filter rating data frame with minimum number of ratings for user/item is usually useful to\n",
    "    generate a new data frame with warm user/item. The warmth is defined by min_rating argument. For\n",
    "    example, a user is called warm if he has rated at least 4 items.\n",
    "    Args:\n",
    "        data (pandas.DataFrame): DataFrame of user-item tuples. Columns of user and item\n",
    "            should be present in the DataFrame while other columns like rating, \n",
    "            timestamp, etc. can be optional.\n",
    "        min_rating (int): minimum number of ratings for user or item.\n",
    "        filter_by (str): either \"user\" or \"item\", depending on which of the two is to \n",
    "            filter with min_rating.\n",
    "        col_user (str): column name of user ID.\n",
    "        col_item (str): column name of item ID.\n",
    "    Returns:\n",
    "        pandas.DataFrame: DataFrame with at least columns of user and item that has been filtered by the given specifications.\n",
    "    \"\"\"\n",
    "    split_by_column = _get_column_name(\n",
    "        filter_by, col_user, col_item\n",
    "    )\n",
    "\n",
    "    if min_rating < 1:\n",
    "        raise ValueError(\"min_rating should be integer and larger than or equal to 1.\")\n",
    "\n",
    "    return (\n",
    "        data\n",
    "        .groupby(split_by_column)\n",
    "        .filter(lambda x: len(x) >= min_rating)\n",
    "    )\n",
    "\n",
    "\n",
    "def _do_stratification(\n",
    "    data,\n",
    "    ratio=0.75,\n",
    "    min_rating=1,\n",
    "    filter_by=\"user\",\n",
    "    is_random=True,\n",
    "    seed=42,\n",
    "    col_user='USERID',\n",
    "    col_item='ITEMID',\n",
    "    col_timestamp='TIMESTAMP',\n",
    "):\n",
    "    # A few preliminary checks.\n",
    "    if not (filter_by == \"user\" or filter_by == \"item\"):\n",
    "        raise ValueError(\"filter_by should be either 'user' or 'item'.\")\n",
    "\n",
    "    if min_rating < 1:\n",
    "        raise ValueError(\"min_rating should be integer and larger than or equal to 1.\")\n",
    "\n",
    "    if col_user not in data.columns:\n",
    "        raise ValueError(\"Schema of data not valid. Missing User Col\")\n",
    "\n",
    "    if col_item not in data.columns:\n",
    "        raise ValueError(\"Schema of data not valid. Missing Item Col\")\n",
    "\n",
    "    if not is_random:\n",
    "        if col_timestamp not in data.columns:\n",
    "            raise ValueError(\"Schema of data not valid. Missing Timestamp Col\")\n",
    "\n",
    "    multi_split, ratio = process_split_ratio(ratio)\n",
    "\n",
    "    split_by_column = col_user if filter_by == \"user\" else col_item\n",
    "\n",
    "    ratio = ratio if multi_split else [ratio, 1 - ratio]\n",
    "\n",
    "    if min_rating > 1:\n",
    "        data = min_rating_filter_pandas(\n",
    "            data,\n",
    "            min_rating=min_rating,\n",
    "            filter_by=filter_by,\n",
    "            col_user=col_user,\n",
    "            col_item=col_item,\n",
    "        )\n",
    "\n",
    "    # Split by each group and aggregate splits together.\n",
    "    splits = []\n",
    "\n",
    "    # If it is for chronological splitting, the split will be performed in a random way.\n",
    "    df_grouped = (\n",
    "        data.sort_values(col_timestamp).groupby(split_by_column)\n",
    "        if is_random is False\n",
    "        else data.groupby(split_by_column)\n",
    "    )\n",
    "\n",
    "    for _, group in df_grouped:\n",
    "        group_splits = split_pandas_data_with_ratios(\n",
    "            group, ratio, shuffle=is_random, seed=seed\n",
    "        )\n",
    "\n",
    "        # Concatenate the list of split dataframes.\n",
    "        concat_group_splits = pd.concat(group_splits)\n",
    "\n",
    "        splits.append(concat_group_splits)\n",
    "\n",
    "    # Concatenate splits for all the groups together.\n",
    "    splits_all = pd.concat(splits)\n",
    "\n",
    "    # Take split by split_index\n",
    "    splits_list = [\n",
    "        splits_all[splits_all[\"split_index\"] == x].drop(\"split_index\", axis=1)\n",
    "        for x in range(len(ratio))\n",
    "    ]\n",
    "\n",
    "    return splits_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random Split v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def random_split_v2(data, ratio=0.8, seed=42):\n",
    "    \"\"\"Pandas random splitter.\n",
    "    The splitter randomly splits the input data.\n",
    "    Args:\n",
    "        data (pandas.DataFrame): Pandas DataFrame to be split.\n",
    "        ratio (float or list): Ratio for splitting data. If it is a single float number\n",
    "            it splits data into two halves and the ratio argument indicates the ratio\n",
    "            of training data set; if it is a list of float numbers, the splitter splits\n",
    "            data into several portions corresponding to the split ratios. If a list is\n",
    "            provided and the ratios are not summed to 1, they will be normalized.\n",
    "        seed (int): Seed.\n",
    "    Returns:\n",
    "        list: Splits of the input data as pandas.DataFrame.\n",
    "    \"\"\"\n",
    "    multi_split, ratio = process_split_ratio(ratio)\n",
    "\n",
    "    if multi_split:\n",
    "        splits = split_pandas_data_with_ratios(data, ratio, shuffle=True, seed=seed)\n",
    "        splits_new = [x.drop(\"split_index\", axis=1) for x in splits]\n",
    "\n",
    "        return splits_new\n",
    "    else:\n",
    "        return train_test_split(data, test_size=None, train_size=ratio, random_state=seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chrono Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def chrono_split(\n",
    "    data,\n",
    "    ratio=0.75,\n",
    "    min_rating=1,\n",
    "    filter_by=\"user\",\n",
    "    col_user='USERID',\n",
    "    col_item='ITEMID',\n",
    "    col_timestamp='TIMESTAMP',\n",
    "):\n",
    "    \"\"\"Pandas chronological splitter.\n",
    "    This function splits data in a chronological manner. That is, for each user / item, the\n",
    "    split function takes proportions of ratings which is specified by the split ratio(s).\n",
    "    The split is stratified.\n",
    "    Args:\n",
    "        data (pandas.DataFrame): Pandas DataFrame to be split.\n",
    "        ratio (float or list): Ratio for splitting data. If it is a single float number\n",
    "            it splits data into two halves and the ratio argument indicates the ratio of\n",
    "            training data set; if it is a list of float numbers, the splitter splits\n",
    "            data into several portions corresponding to the split ratios. If a list is\n",
    "            provided and the ratios are not summed to 1, they will be normalized.\n",
    "        seed (int): Seed.\n",
    "        min_rating (int): minimum number of ratings for user or item.\n",
    "        filter_by (str): either \"user\" or \"item\", depending on which of the two is to\n",
    "            filter with min_rating.\n",
    "        col_user (str): column name of user IDs.\n",
    "        col_item (str): column name of item IDs.\n",
    "        col_timestamp (str): column name of timestamps.\n",
    "    Returns:\n",
    "        list: Splits of the input data as pandas.DataFrame.\n",
    "    \"\"\"\n",
    "    return _do_stratification(\n",
    "        data,\n",
    "        ratio=ratio,\n",
    "        min_rating=min_rating,\n",
    "        filter_by=filter_by,\n",
    "        col_user=col_user,\n",
    "        col_item=col_item,\n",
    "        col_timestamp=col_timestamp,\n",
    "        is_random=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stratified Split v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def stratified_split_v2(\n",
    "    data,\n",
    "    ratio=0.75,\n",
    "    min_rating=1,\n",
    "    filter_by=\"user\",\n",
    "    col_user='USERID',\n",
    "    col_item='ITEMID',\n",
    "    seed=42,\n",
    "):\n",
    "    \"\"\"Pandas stratified splitter.\n",
    "    For each user / item, the split function takes proportions of ratings which is\n",
    "    specified by the split ratio(s). The split is stratified.\n",
    "    Args:\n",
    "        data (pandas.DataFrame): Pandas DataFrame to be split.\n",
    "        ratio (float or list): Ratio for splitting data. If it is a single float number\n",
    "            it splits data into two halves and the ratio argument indicates the ratio of\n",
    "            training data set; if it is a list of float numbers, the splitter splits\n",
    "            data into several portions corresponding to the split ratios. If a list is\n",
    "            provided and the ratios are not summed to 1, they will be normalized.\n",
    "        seed (int): Seed.\n",
    "        min_rating (int): minimum number of ratings for user or item.\n",
    "        filter_by (str): either \"user\" or \"item\", depending on which of the two is to\n",
    "            filter with min_rating.\n",
    "        col_user (str): column name of user IDs.\n",
    "        col_item (str): column name of item IDs.\n",
    "    Returns:\n",
    "        list: Splits of the input data as pandas.DataFrame.\n",
    "    \"\"\"\n",
    "    return _do_stratification(\n",
    "        data,\n",
    "        ratio=ratio,\n",
    "        min_rating=min_rating,\n",
    "        filter_by=filter_by,\n",
    "        col_user=col_user,\n",
    "        col_item=col_item,\n",
    "        is_random=True,\n",
    "        seed=seed,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-24 07:50:04\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "IPython: 5.5.0\n",
      "numpy  : 1.19.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
