{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Introduction\n",
    "\n",
    "The model package utilizes a series of four modules to prepare data and deploy models to perform machine learning guided directed evolution.\n",
    "\n",
    "1. ```Splitters``` are used to split the dataset into training, validation, and test sets.\n",
    "\n",
    "2. ```Featurizers``` are used to featurize the sequences.\n",
    "\n",
    "3. ```Predictors``` are used to train and deploy models to perform property prediction.\n",
    "\n",
    "4. ```Proposers``` are used to propose and evaluate new sequences given a list of trained models.\n",
    "\n",
    "For developers, additional classes can be added to each module to implement custom functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from model.splitters import *\n",
    "from model.featurizers import *\n",
    "from model.predictors import *\n",
    "from model.proposers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "First, define the following variables:\n",
    "- ```experiment_name```: the name of the experiment\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "experiment_name = \"example_experiment\"\n",
    "protein_name = \"example_protein\"\n",
    "wt_file = \"../../../data/example_protein/apex.fasta\"\n",
    "training_dataset_fname = '../../../data/example_protein/example_dataset.csv'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Datasets should be in CSV format with following columns:\n",
    "\n",
    "- ```mutation```: the mutation, formatted as ```A123V```, wherein multi-mutants are separated by forward slashes (```/```). If there is no mutation, the value should be ```WT```.\n",
    "\n",
    "- ```property_value```: the property value\n",
    "\n",
    "- ```evolution_round```: the evolution round in which the variant was measured (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv(training_dataset_fname)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Splitters\n",
    "\n",
    "Several splitters are available in the ```splitters``` module. Each splitter can split the dataset into training, validation, and test sets using different strategies. To learn more about the splitters, check out the ```splitters.ipynb``` notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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.\n",
    "\n",
    "We will initilize two splitters: one for non-neural network models and one for neural network models. We will use a validation set of 15% of the data for the neural network models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "splitter = RandomProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=True, y_scaling=False, val_split=None)\n",
    "splitter_nn = RandomProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=True, y_scaling=False, val_split=0.15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```data``` attribute of the splitter views the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "splitter.data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After initializing the splitter, run ```splitter.split_data()``` to split the data. For ```RandomProteinSplitter```, the ```split_data()``` method takes the following parameters:\n",
    "\n",
    "- ```test_size```: the proportion of the dataset to include in the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "splitter.split_data(test_size=0.15)\n",
    "splitter_nn.split_data(test_size=0.15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it! The dataset has now been split into training, validation, and test sets and can be fed into the ```Predictors``` module to train and deploy models. If you check the ```splits``` attribute of the splitter, 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": [
    "splitter.splits.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(splitter.splits['X_train'][:3])\n",
    "print(splitter.splits['y_train'][:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Featurizers\n",
    "\n",
    "Featurizers are used to featurize the sequences. To learn more about the different featurizers, check out the ```featurizers.ipynb``` notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each featurizer has the following parameters:\n",
    "\n",
    "- ```protein```: the name of the protein for caching\n",
    "\n",
    "- ```use_cache```: whether to cache the features for later use (default: ```False```)\n",
    "\n",
    "- ```flatten_features```: whether to flatten the feature vectors (default: ```False```)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "featurizer = OneHotFeaturizer(protein=protein_name, use_cache=True, flatten_features=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predictors\n",
    "\n",
    "Predictors are used to train and deploy models to perform property prediction. To learn more about the different predictors, check out the documentation.\n",
    "\n",
    "Each predictor has the following parameters:\n",
    "\n",
    "- ```splitter```: the splitter to use\n",
    "\n",
    "- ```featurizer```: the featurizer to use\n",
    "\n",
    "- ```use_cache```: whether to cache the model for later use (default: ```False```)\n",
    "\n",
    "- ```show_plots```: whether to show matplotlib plots (default: ```True```)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training non-neural network models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There a several models available in the ```predictors``` module. To learn more about the different models, check out the documentation.\n",
    "\n",
    "- ```RidgeRegressor```: a ridge regression model\n",
    "\n",
    "- ```RandomForestRegressor```: a random forest regression model\n",
    "\n",
    "- ```GPLinearRegressor```: a gaussian process linear regression model\n",
    "\n",
    "- ```GPQuadRegressor```: a gaussian process quadratic regression model\n",
    "\n",
    "- ```GPRBFRegressor```: a gaussian process radial basis function regression model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictor = RidgeRegressor(splitter, featurizer, use_cache=True, show_plots=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After initializing the predictor, run ```predictor.run_model()``` to train and deploy the model. This command returns a dictionary of performance statistics as well as a plot of the model's performance on the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = predictor.run_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame([stats]).transpose().rename(columns={0: 'Value'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! The model has now been trained and deployed and can be used to predict the property of new sequences."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training neural network models\n",
    "\n",
    "Training neural network models is similar to training machine learning models. The only major difference is that neural network models require a ```config``` dictionary to specify the network architecture.\n",
    "\n",
    "In the model package, there are two simple neural network models available: ```Fcn``` and ```Cnn```. ```Fcn``` is a fully connected neural network and ```Cnn``` is a convolutional neural network.\n",
    "\n",
    "First, we will train a fully connected neural network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fully connected neural network\n",
    "\n",
    "The config dictionary for the fully connected neural network has the following parameters:\n",
    "\n",
    "- ```layer_size```: the number of neurons in the hidden layers\n",
    "\n",
    "- ```num_layers```: the number of hidden layers\n",
    "\n",
    "- ```learning_rate```: the learning rate for the optimizer\n",
    "\n",
    "- ```batch_size```: the batch size for training\n",
    "\n",
    "- ```optimizer```: the optimizer to use (default: ```adam```)\n",
    "\n",
    "- ```epochs```: the number of epochs to train for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# config\n",
    "config = {\n",
    "          'layer_size': 100,\n",
    "          'num_layers' : 2,\n",
    "          'learning_rate': 0.001,\n",
    "          'batch_size': 32,\n",
    "          'optimizer': 'adam',\n",
    "          'epochs': 300\n",
    "}\n",
    "\n",
    "fcn_model = Fcn(splitter_nn, featurizer, config=config, use_cache=True, show_plots=True)\n",
    "stats = fcn_model.run_model()\n",
    "pd.DataFrame([stats]).transpose().rename(columns={0: 'Value'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convolutional neural network\n",
    "\n",
    "The convolutional neural network is a 2D convolutional neural network that scans across the featurized protein sequences with dimensions of ```(sequence_length, feature_length)``` with filter size ```(kernel_size, feature_length)```. When using the ```Cnn``` class, make sure to set ```flatten_features=False``` in the ```Featurizer``` class.\n",
    "\n",
    "The config dictionary for the convolutional neural network has the following parameters:\n",
    "\n",
    "- ```layersize_filtersize```: the number of hidden layers and the number of filters separated by a dash (```-```)\n",
    "\n",
    "- ```kernel_size```: the kernel size for the convolutional layer\n",
    "\n",
    "- ```learning_rate```: the learning rate for the optimizer\n",
    "\n",
    "- ```batch_size```: the batch size for training\n",
    "\n",
    "- ```optimizer```: the optimizer to use (default: ```adam```)\n",
    "\n",
    "- ```epochs```: the number of epochs to train for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "          'layersize_filtersize': \"1-12\",\n",
    "          'kernel_size' : 17,\n",
    "          'learning_rate':0.001,\n",
    "          'batch_size': 32,\n",
    "          'optimizer': 'adam',\n",
    "          'epochs': 10\n",
    "}\n",
    "\n",
    "cnn_model = Cnn(splitter_nn, featurizer, config=config, use_cache=True)\n",
    "stats = cnn_model.run_model()\n",
    "pd.DataFrame([stats]).transpose().rename(columns={0: 'Value'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Proposers\n",
    "\n",
    "Proposers are used to propose new sequences given a list of trained models. To learn more about the different proposers, check out the documentation. Generally, we used the ```CombinatorialProposer``` to propose new sequences.\n",
    "\n",
    "Each proposer has the following parameters:\n",
    "\n",
    "- ```start_seq```: the starting sequence to mutate, generally the wildtype sequence\n",
    "\n",
    "- ```models```: the list of trained models\n",
    "\n",
    "- ```trust_radius```: the maximum number of mutations allowed in the proposed variant\n",
    "\n",
    "- ```num_seeds```: the maximum number of sequences to propose for evaluation, -1 means tests all possible variants\n",
    "\n",
    "- ```mutation_pool```: the list of allowed mutations for generating the proposed variants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "wt_seq = 'MGKSYPTVSADYQDAVEKAKKKLRGFIAEKRCAPLMLRLAFHSAGTFDKGTKTGGPFGTIKHPAELAHSANNGLDIAVRLLEPLKAEFPILSYADFYQLAGVVAVEVTGGPKVPFHPGREDKPEPPPEGRLPDATKGSDHLRDVFGKAMGLTDQDIVALSGGHTIGAAHKERSGFEGPWTSNPLIFDNSYFTELLSGEKEGLLQLPSDKALLSDPVFRPLVDKYAADEDAFFADYAEAHQKLSELGFADA'\n",
    "mutations = ['T192V', 'T192K', 'A167R', 'N72A', 'D222E', 'A148Q', 'D229A', 'S138A', 'K61R', 'S196A', 'I185V', 'L84V', 'E87Q', 'G50R', 'L80M']\n",
    "\n",
    "proposer = CombinatorialProposer(\n",
    "    start_seq=wt_seq,\n",
    "    models=[fcn_model], \n",
    "    trust_radius=10, \n",
    "    num_seeds=-1, \n",
    "    # num_seeds=20, \n",
    "    mutation_pool=mutations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After initializing the proposer, run ```proposer.propose()``` to propose new sequences. This command returns a dataframe of proposed sequences and their evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proposal_results = proposer.propose()\n",
    "proposal_results.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After proposing new sequences, run ```proposer.evaluate_proposals()``` to evaluate the proposed sequences. This command returns a dataframe of proposed sequences and their evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proposer.evaluate_proposals()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```proposals``` dataframe now contains the proposed sequences and their evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proposer.proposals.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results can now be saved to a CSV file using ```proposer.save_proposals()```. Results will be saved in the following folder: ```destination_folder/proposers/results/```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proposer.save_proposals(f'{experiment_name}_proposals') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another alternative proposer is the ```DeepMutationalScanningProposer```, which generates every possible single amino acid substitution and predicts the property of each proposed sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dms_proposer = DeepMutationalScanningProposer(\n",
    "    start_seq=wt_seq, \n",
    "    models=[predictor]\n",
    "    )\n",
    "dms_proposer.propose()\n",
    "dms_proposer.evaluate_proposals()\n",
    "dms_proposer.save_proposals(f'{experiment_name}_dms_proposals') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dms_proposer.proposals.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Aside) Cache save locations\n",
    "\n",
    "A new directory named ```proteins``` will be created. Under the ```protein_name```, there will be cache folders for splitters, featurizers, predictors, and proposers. The cache folder organization will look like this:\n",
    "\n",
    "```\n",
    "example_protein/         \n",
    "鈹溾攢鈹€ example_dataset.csv  \n",
    "鈹溾攢鈹€ feature_cache/       \n",
    "鈹?  鈹斺攢鈹€ onehot/          \n",
    "鈹溾攢鈹€ model_cache/         \n",
    "鈹?  鈹斺攢鈹€ example_dataset/ \n",
    "鈹?      鈹溾攢鈹€ objects/     \n",
    "鈹?      鈹斺攢鈹€ results/     \n",
    "鈹斺攢鈹€ proposers/      \n",
    "    鈹斺攢鈹€ results/         \n",
    "鈹溾攢鈹€ split_cache/         \n",
    "鈹?  鈹斺攢鈹€ example_dataset/ \n",
    "```\n",
    "\n",
    "The ```feature_cache``` folder contains the featurized sequences separated based on featurizer type.\n",
    "\n",
    "The ```model_cache``` folder contains the predictor objects separated by dataset. The ```objects``` folder contains the saved models and the ```results``` folder contain results generating when comparing multiple models (seen later in the ```Part_2_comparing_multiple_models.ipynb``` notebook).\n",
    "\n",
    "The ```proposers``` folder contains the results of the evaluated proposed sequences.\n",
    "\n",
    "The ```split_cache``` folder contains the splitter objects separated by dataset. \n",
    "\n",
    "If you check the ```file_attrs``` attribute of the splitter or predictor, you will see the cache save locations of the objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "splitter.file_attrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.file_attrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "Overall, we have seen how to use the ```splitters```, ```featurizers```, ```predictors```, and ```proposers``` modules to train and deploy models to perform property prediction in a few lines of code. We separated the code into these four modules to be able to compare different methods of data splits, featurizations, and models. The full example of a code block training a simple ridge regression model and proposing new sequences is shown below. By changing a single line of code, you can test a different data split method, featurization, or model, allowing for easy comparison of different methods.\n",
    "\n",
    "For streamlined comparison of multiple methods of data splits, featurizations, and models, head over to ```Part_2_comparing_multiple_models.ipynb```.\n",
    "\n",
    "Again, to learn more about the different modules, check out the ```splitters.ipynb``` and ```featurizers.ipynb``` notebooks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Full Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define variables\n",
    "experiment_name = \"example_experiment\"\n",
    "protein_name = \"example_protein\"\n",
    "wt_file = \"../../../data/example_protein/apex.fasta\"\n",
    "training_dataset_fname = '../../../data/example_protein/example_dataset.csv'\n",
    "\n",
    "# Initialize splitter\n",
    "splitter = RandomProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=True, y_scaling=False, val_split=None)\n",
    "splitter.split_data(test_size=0.15)\n",
    "\n",
    "# Initialize featurizer\n",
    "featurizer = OneHotFeaturizer(protein=protein_name, use_cache=True, flatten_features=False)\n",
    "\n",
    "# Initialize predictor\n",
    "predictor = RidgeRegressor(splitter, featurizer, use_cache=True)\n",
    "stats = predictor.run_model()\n",
    "\n",
    "# Initialize proposer\n",
    "wt_seq = 'MGKSYPTVSADYQDAVEKAKKKLRGFIAEKRCAPLMLRLAFHSAGTFDKGTKTGGPFGTIKHPAELAHSANNGLDIAVRLLEPLKAEFPILSYADFYQLAGVVAVEVTGGPKVPFHPGREDKPEPPPEGRLPDATKGSDHLRDVFGKAMGLTDQDIVALSGGHTIGAAHKERSGFEGPWTSNPLIFDNSYFTELLSGEKEGLLQLPSDKALLSDPVFRPLVDKYAADEDAFFADYAEAHQKLSELGFADA'\n",
    "mutations = ['T192V', 'T192K', 'A167R', 'N72A', 'D222E', 'A148Q', 'D229A', 'S138A', 'K61R', 'S196A', 'I185V', 'L84V', 'E87Q', 'G50R', 'L80M']\n",
    "proposer = CombinatorialProposer(start_seq=wt_seq, models=[predictor], trust_radius=10, num_seeds=20, mutation_pool=mutations)\n",
    "proposal_results = proposer.propose()\n",
    "proposer.evaluate_proposals()\n",
    "proposer.save_proposals(f'{experiment_name}_proposals') "
   ]
  }
 ],
 "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_mac",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}