{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This notebook will discuss how to use the various splitters in the ```multievolve``` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from model.splitters import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "First, define the following variables, including a structure of the protein of interest:\n",
    "\n",
    "- ```protein_name```: the name of the protein\n",
    "\n",
    "- ```wt_file```: the path to the wildtype sequence\n",
    "\n",
    "- ```training_dataset_fname```: the path to the training dataset\n",
    "\n",
    "- ```structure_file```: the path to the structure file, either .pdb or .cif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "protein_name = \"example_protein\"\n",
    "wt_file = \"../../../data/example_protein/apex.fasta\"\n",
    "training_dataset_fname = '../../../data/example_protein/example_dataset.csv'\n",
    "structure_file = \"../../../data/example_protein/apex.cif\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Refresher\n",
    "\n",
    "As previously mentioned, each splitter has the following parameters:\n",
    "\n",
    "- ```protein_name```: the name of the protein\n",
    "\n",
    "- ```training_dataset_fname```: the path to the training dataset\n",
    "\n",
    "- ```wt_file```: the path to the wildtype sequence\n",
    "\n",
    "- ```csv_has_header```: whether the CSV has a header\n",
    "\n",
    "- ```use_cache```: whether to cache the processed dataset for later use (default: ```False```)\n",
    "\n",
    "- ```y_scaling```: whether to scale the property values between 0 and 1 (default: ```False```)\n",
    "\n",
    "- ```val_split```: the proportion of the dataset to include in the validation set (default: ```None```). The validation set is only used for when training neural network models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### KFoldProteinSplitter\n",
    "\n",
    "```KFoldProteinSplitter```: Performs k-fold cross-validation by randomly splitting data into k folds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "kfold_splitter = KFoldProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike the other Splitters in which we run ```split_data()``` method, we obtain the processed datasets by running ```kfold_splitter.generate_splits(n_splits=5)```, where ```n_splits``` is the number of folds, in this case we perform 5-fold cross-validation. \n",
    "\n",
    "This returns a list of ```n_splits``` splitter objects, each with using a different fold for the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "splits = kfold_splitter.generate_splits(n_splits=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, if you check the ```splits``` attribute of one of the splitter objects, then you will see that the dataset has been split into training, validation, and test sets in the form of a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "splits[0].splits.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RoundProteinSplitter\n",
    "\n",
    "```RoundProteinSplitter```: Splits data based on evolution rounds, allowing training on early rounds and testing on later rounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "round_splitter = RoundProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```RoundProteinSplitter```,```split_data()``` has the following arguments:\n",
    "- ```max_train_round```: the maximum round number to include in the training set\n",
    "- ```min_test_round```: the minimum round number to include in the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "round_splitter.split_data(max_train_round=0, min_test_round=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RandomProteinSplitter\n",
    "\n",
    "```RandomProteinSplitter```: Randomly splits data into training and test sets with a specified test size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "random_splitter = RandomProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```RandomProteinSplitter```,```split_data()``` has the following arguments:\n",
    "- ```test_size```: the proportion of the dataset to include in the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "random_splitter.split_data(test_size=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PositionProteinSplitter\n",
    "\n",
    "```PositionProteinSplitter```: Splits based on mutation positions - variants with mutations at certain positions go to test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "position_splitter = PositionProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```PositionProteinSplitter```,```split_data()``` has the following arguments:\n",
    "- ```test_size_sample```: the proportion of the dataset to sample to get mutation positions to exclude out of the training set\n",
    "- ```sample_iter```: the number of iterations to perform to get a test set size between ```test_size_min``` and ```test_size_max```\n",
    "- ```test_size_min```: the minimum test size set desired\n",
    "- ```test_size_max```: the maximum test size set allowed\n",
    "\n",
    "When splitting the data, the splitter will sample random mutations to get the mutation positions to exclude out of the training set. The splitter will attempt to get a test set size between a specified range, and will repeat sampling for specified number of iterations if the test set size is not within the desired range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "position_splitter.split_data(test_size_sample=0.2, sample_iter=3, test_size_min=0.2, test_size_max=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RegionProteinSplitter\n",
    "\n",
    "```RegionProteinSplitter```: Splits based on protein regions - variants with mutations in specified regions go to test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "region_splitter = RegionProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```RegionProteinSplitter```,```split_data()``` has the following arguments:\n",
    "- ```region```: a list of two numbers defining the minimum and maximum positions to include in the test set (e.g. [1, 60])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "region_splitter.split_data(region=[1, 60])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PropertyProteinSplitter\n",
    "\n",
    "```PropertyProteinSplitter```: Splits based on property values - can separate high/low performing variants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "property_splitter = PropertyProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```PropertyProteinSplitter```,```split_data()``` has the following arguments:\n",
    "- ```property```: the value of the property split on\n",
    "- ```above_or_below```: 'above' or 'below', values to leave out into the test set based on the given property value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "property_splitter.split_data(property=1, above_or_below='above')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MutLoadProteinSplitter\n",
    "\n",
    "```MutLoadProteinSplitter```: Splits based on number of mutations - can train on low mutation count variants and test on higher ones.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mutload_splitter = MutLoadProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```MutLoadProteinSplitter```,```split_data()``` has the following arguments:\n",
    "- ```max_train_muts```: the maximum mutation load to include in the training set\n",
    "- ```min_test_muts```: the minimum mutation load to include in the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mutload_splitter.split_data(max_train_muts=1, min_test_muts=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ResidueDistanceSplitter\n",
    "\n",
    "```ResidueDistanceSplitter```: Splits based on 3D distances between mutations using protein structure.\n",
    "\n",
    "When initializing ```ResidueDistanceSplitter```, we need to specify the additional arguments:\n",
    "- ```pdb_file```: the path to the PDB/CIF structure file\n",
    "- ```chain_ids```: the chain IDs of the protein of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "residue_distance_splitter = ResidueDistanceSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=False, y_scaling=False, val_split=None,\n",
    "                                                    pdb_file=structure_file, chain_ids=['A'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ```ResidueDistanceSplitter```,```split_data()``` has the following arguments:\n",
    "- ```percentile_threshold```: the percentile threshold for the distance to include in the training set\n",
    "- ```min_test_muts```: the minimum number of mutations to include in the test set\n",
    "- ```max_train_muts```: the maximum number of mutations to include in the training set\n",
    "- ```randomized_control```: whether to randomize the distance dictionary as a control (default: ```False```)\n",
    "\n",
    "When splitting the data, the splitter will calculate the distance percentile for each variant within its mutational load group. The splitter will consider variants with mutational load less than or equal to the ```max_train_muts``` for the training set. The splitter will then split the data based on the distance percentile, with variants with distances less than or equal to the percentile threshold going to the training set and variants with distances greater than the percentile threshold going to the test set. All variants will a mutational load higher than or equal to ```max_train_muts``` will go to the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "residue_distance_splitter.split_data(\n",
    "        percentile_threshold=50,  \n",
    "        min_test_muts=5,         \n",
    "        max_train_muts=2,\n",
    "        randomized_control=False\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-chain proteins\n",
    "\n",
    "If you are working with multi-chain proteins such as antibodies that have a heavy variable domain and light variable domain, you can use the Splitters to accept both chains.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "With multi-chain proteins, you need to:\n",
    "- Specify the wild-type sequences for each chain in the ```wt_files``` argument as a list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "protein_name_multichain = \"example_multichain_protein\"\n",
    "wt_files = ['../../../data/example_multichain_protein/vh_chain1.fasta', '../../../data/example_multichain_protein/vl_chain2.fasta']\n",
    "training_dataset_fname_multichain = '../../../data/example_multichain_protein/example_dataset.csv'\n",
    "structure_file_multichain = 'multichain_protein.cif'\n",
    "chain_ids = ['A', 'B']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formatting datasets for multi-chain proteins\n",
    "For multi-chain datasets, the mutation strings for each chain are separated by a colon (e.g. ```F32Y:E61Y```). If a variant is wild-type for both chains, then the mutation string should be ```WT:WT```. If a variant is wild-type for one chain and has a mutation for the other chain, then the mutation string should be ```WT:F32Y```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(training_dataset_fname_multichain)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example with ResidueDistanceSplitter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the ```ResidueDistanceSplitter```, the arguments should be as follows:\n",
    "- Specify the chain IDs in the ```chain_ids``` argument as a list.\n",
    "- Specify the structure file in the ```pdb_file``` argument. This should be one structure containing all chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = ResidueDistanceSplitter(protein_name_multichain, training_dataset_fname_multichain, wt_files, csv_has_header=True, use_cache=False, \n",
    "                 y_scaling=False,\n",
    "                 val_split=None,\n",
    "                 pdb_file=structure_file_multichain,\n",
    "                 chain_ids=chain_ids,\n",
    "                 random_state=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For multi-chain proteins, the Splitter will concatenate the sequences of each chain to get the full sequence. It will then automatically adjust the mutation positions for each chain to match the positions in the full concatenatedsequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "residue_distance_splitter.split_data(\n",
    "        percentile_threshold=50,  \n",
    "        min_test_muts=5,         \n",
    "        max_train_muts=2,\n",
    "        randomized_control=False\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "multievolve",
   "name": "workbench-notebooks.m128",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m128"
  },
  "kernelspec": {
   "display_name": "multievolve",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}