{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Splitting and subsetting PPIs\n",
    "\n",
    "The PPIRef package provides a unified approach to storing and processing data splits and other subsets of PPIs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ppiref.split import read_split, read_fold, read_split_source, write_split\n",
    "from ppiref.definitions import PPIREF_TEST_DATA_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example 1.** Create a toy data split. In this example, we use some of PPIs extracted in the \"Extracting PPIs\" tutorial."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, prepare PPI ids for the split and use the `write_split` function to store the split as a JSON file. This will run simple sanity checks on the split (e.g. that PPI ids are not overlapping across the folds or the split is complete with respect to the source directory with PPI files). In this example, we introduce a data leakage on purpose by including the same PPI in both the training and the test sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/anton/dev/PPIRef/ppiref/split.py:63: UserWarning: Folds train and test are not disjoint.\n",
      "  warnings.warn(f'Folds {fold_a} and {fold_b} are not disjoint.')\n",
      "/Users/anton/dev/PPIRef/ppiref/split.py:69: UserWarning: Split is not complete: 2 of 38 PPIs contained.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "split = {'train': ['10gs_A_B'], 'test': ['1ahw_B_C', '10gs_A_B']}\n",
    "write_split('demo_split', source=PPIREF_TEST_DATA_DIR / 'ppi_dir', folds=split)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you can read the split using the `read_split` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                              \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'train': [PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/0g/10gs_A_B.pdb')],\n",
       " 'test': [PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/ah/1ahw_B_C.pdb'),\n",
       "  PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/0g/10gs_A_B.pdb')]}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "read_split('demo_split')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or read individual folds using the `read_fold` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                              \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/0g/10gs_A_B.pdb')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "read_fold('demo_split', 'train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                              \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/0g/10gs_A_B.pdb'),\n",
       " PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/ah/1ahw_B_C.pdb'),\n",
       " PPIPath('/Users/anton/dev/PPIRef/ppiref/data/test/ppi_dir/0g/10gs_A_B.pdb')]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "read_fold('demo_split', 'train+test')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example 2.** Read pre-defined data splits."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PPIRef50K used to train PPIformer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['4q2p_A_B', '3q2s_A_D', '6q2v_B_E']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fold = read_fold('ppiref_10A_filtered_clustered_03', 'whole', full_paths=False)\n",
    "fold[:3]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test set from non-leaking SKEMPI v2.0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1B3S_A_D', '1B2U_A_D', '1BRS_A_D']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fold = read_fold('skempi2_iclr24_split', 'test', full_paths=False)\n",
    "fold[:3]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DIPS set used to train and validate EquiDock and DiffDock-PP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1v6j_A_D', '2v6a_A_L', '2v6a_B_O']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fold = read_fold('dips_equidock', 'train+val', full_paths=False)\n",
    "fold[:3]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each of the pre-defined splits, you can read the source directory of the PPIs using the `read_split_source` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/Users/anton/dev/PPIRef/ppiref/data/ppiref/ppi_6A')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "read_split_source('ppiref_6A_raw')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ppiref",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}