{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # (optional) check if the installation is successful\n",
    "# import torch_sparse\n",
    "# import torch_geometric\n",
    "# import torch_cluster\n",
    "# import torch_scatter\n",
    "# import torch_spline_conv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This notebook shows you how to run zero-shot with an ensemble of non-structure-informed models (ESM-1v, ESM-2 3B) with the ```zero_shot_esm_dms``` function as well as run zero-shot with a structure-informed model (ESM-IF) using ```zero_shot_esm_if_dms``` function. \n",
    "\n",
    "```zero_shot_esm_dms``` requires the wild-type amino acid sequence of the protein of the interest\n",
    "\n",
    "```zero_shot_esm_if_dms``` requires the wild-type amino acid sequence and the structure of the protein of the interest, including the chain id of the protein of interest in the structure file.\n",
    "\n",
    "Both functions return a dataframe with all the possible single amino acid mutations and their corresponding log likelihood ratio scores (i.e. the ratio of the likelihood compared to the wild-type sequence)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from Bio import SeqIO\n",
    "\n",
    "from model import zero_shot_esm_dms, zero_shot_esm_if_dms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "wt_file = \"../../../data/example_protein/apex.fasta\"\n",
    "pdb_file = \"../../../data/example_protein/apex.cif\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "wt_seq = str(SeqIO.read(wt_file, \"fasta\").seq)\n",
    "\n",
    "esm_zeroshot = zero_shot_esm_dms(wt_seq)\n",
    "esm_if_zeroshot = zero_shot_esm_if_dms(wt_seq, pdb_file, chain_id = 'A', scoring_strategy='wt-marginals')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```zero_shot_esm_dms``` returns a dataframe with the log likelihood ratio scores for each model as well as whether the mutation had a ratio greater than 1 indicated by the corresponding ```model#_pass``` column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "esm_zeroshot.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```zero_shot_esm_if_dms``` returns a dataframe with the log likelihood ratio scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "esm_if_zeroshot.head(10)"
   ]
  }
 ],
 "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_public",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}