{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This notebook will discuss how to use the various featurizers in the ```multievolve``` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.splitters import *\n",
    "from model.featurizers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "First, define the following variables:\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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": [
    "Define a splitter object 鈥?we will just use this to load the dataset and pull sequences from to featurizer later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "splitter = RandomProteinSplitter(protein_name, training_dataset_fname, wt_file, csv_has_header=True, use_cache=True, y_scaling=False, val_split=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Featurizers\n",
    "\n",
    "There are many featurizers available in the ```multievolve``` package. We discuss a few of the most common ones below. \n",
    "\n",
    "- ```OneHotFeaturizer```: One-hot encoding of the protein sequence\n",
    "\n",
    "- ```GeorgievFeaturizer```: Georgiev et al. (2022) featurizer\n",
    "\n",
    "- ```AAIdxFeaturizer```: amino acid index featurizer    \n",
    "\n",
    "- ```ESMLogitsFeaturizer```: ESM-2 logits featurizer\n",
    "\n",
    "- ```ESM2EmbedFeaturizer```: ESM-2 embedding featurizer\n",
    "\n",
    "There are also combinatorial featurizers that combine multiple featurizers.\n",
    "\n",
    "- ```ESMAugmentedFeaturizer```: One-hot encoding augmented likelihood scores from the ESM-1/ESM-2 models\n",
    "\n",
    "- ```OnehotAndGeorgievFeaturizer```: One-hot encoding combined with Georgiev et al. (2022) featurizer, wherein the encodings are stacked along the last axis (i.e. by position)\n",
    "\n",
    "- ```OnehotAndAAIdxFeaturizer```: One-hot encoding augmented with amino acid index featurizer, wherein the encodings are stacked along the last axis (i.e. by position)\n",
    "\n",
    "- ```OnehotAndESMLogitsFeaturizer```: One-hot encoding augmented with ESM-2 logits featurizer, wherein the encodings are stacked along the last axis (i.e. by position)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Base Featurizers\n",
    "onehot = OneHotFeaturizer(protein=protein_name, use_cache=True)\n",
    "georgiev = GeorgievFeaturizer(protein=protein_name, use_cache=True)\n",
    "aa_idx = AAIdxFeaturizer(protein=protein_name, use_cache=True)\n",
    "esm_logits = ESMLogitsFeaturizer(protein=protein_name, use_cache=True)\n",
    "esm_embed = ESM2EmbedFeaturizer(protein=protein_name, use_cache=True)\n",
    "\n",
    "# Combinatorial Featurizers\n",
    "esm_augmented = ESMAugmentedFeaturizer(protein=protein_name, use_cache=True, wt_file=wt_file)\n",
    "onehotgeorgiev = OnehotAndGeorgievFeaturizer(protein=protein_name, use_cache=True)\n",
    "onehotaaidx = OnehotAndAAIdxFeaturizer(protein=protein_name, use_cache=True)\n",
    "onehotesmlogits = OnehotAndESMLogitsFeaturizer(protein=protein_name, use_cache=True)\n",
    "onehotesmmsalogits = OnehotAndESMMSALogitsFeaturizer(protein=protein_name, use_cache=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Featurizers have the function ```featurize```, which takes in a list of sequences and returns the featurized sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_sequences = splitter.data[0][:5].tolist()\n",
    "\n",
    "onehot.featurize(example_sequences)"
   ]
  }
 ],
 "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
}