protein_generator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "38c438d4-7fa4-4a04-b753-25ad6dd2e585",
   "metadata": {},
   "source": [
    "## PROTEIN GENERATOR to test runs and play around"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d728f8d5-6bee-40b4-83b1-a0871c2eb423",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os, sys, json\n",
    "sys.path.append(sys.path[0]+'/utils')\n",
    "sys.path.append(sys.path[0]+'/model')\n",
    "\n",
    "with open('./examples/args.json','r') as f:\n",
    "    args = json.load(f)\n",
    "    \n",
    "# import sampler to run inference    \n",
    "from sampler import SEQDIFF_sampler\n",
    "\n",
    "S = SEQDIFF_sampler()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f4c009e-bed5-489f-8993-a95f7e136aaa",
   "metadata": {},
   "source": [
    "### Specify parameters here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "96fffbb6-f282-4e25-8e96-5d7e7753c35e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model checkpoint...\n",
      "Successfully loaded model checkpoint\n"
     ]
    }
   ],
   "source": [
    "# ARGUMENTS\n",
    "args['sequence'] = 'XXXXXXXXXXXXXXXXPEPSEQXXXXXXXXXXXXXXXX' #None # starting sequence where X is mask token for example: XXXXXXXXXXXPEPSEQXXXXXXXXXXXX\n",
    "args['pdb'] = None #'/home/jgershon/Desktop/SEQDIFF/tmp/test_000000.pdb' # path to pdb to diffuse\n",
    "args['contigs'] = None # contigs for design use, to sample range use ['70-80'] aa, to add strucutral and sequence conditioning from pdb use chain\n",
    "args['out'] = './examples/out/design' # path with prefix to where you want to save outputs\n",
    "args['save_best_plddt'] = True # if true will save best plddt example in the trajectory\n",
    "args['hotspots'] = None # residues to have model target for binder design probelms specified as a list of chain and pdb index example 'B23,B45,B66'\n",
    "args['save_args'] = True # will save args as json\n",
    "args['dump_trb'] = True # will save trb with meta data from run\n",
    "\n",
    "# SECONDARY STRUCTURE BIAS (options a, b, c)\n",
    "# option a\n",
    "args['secondary_structure'] = None # dssp string 'XXXXXHHHHHHHXXXXLLLLXXXXEEEEEEEEEEEEEXXX' for each residue\n",
    "# option b\n",
    "args['dssp_pdb'] = None # path to pdb to copy dssp features of\n",
    "# option c\n",
    "args['helix_bias'] = 0.00 # fraction of sequence to randomly bias toward helix, recommend [0,0.05]\n",
    "args['strand_bias'] = 0.00 # fraction of sequence to randomly bias toward helix, recommend [0,0.05]\n",
    "args['loop_bias'] = 0.00 # fraction of sequence to randomly bias toward helix, recommend [0,0.05]\n",
    "\n",
    "# GUIDING POTENTIALS (can use in combination)\n",
    "# aa_bias\n",
    "args['aa_composition'] = 'W0.2' # amino acid one letter char and fraction of seq to bias as a list 'W0.2,Y0.1'\n",
    "# charge\n",
    "args['target_charge'] = -10 # target charge to bias toward\n",
    "args['target_pH'] = 7.4 # target ph to design in \n",
    "# hydrophobic\n",
    "args['hydrophobic_score'] = -10 # hydrophobic GRAVY score to bias toward (use negative score to bias away from hyrdophobic\n",
    "# for any potential used you must specify the following\n",
    "args['potentials'] = '' # comma seperated list of potentials used 'aa_bias,charge,hydrophobic,'\n",
    "args['potential_scale'] = '' # comma seperated list of weights for each potential in same order as potentials \n",
    "\n",
    "# DIFFUSION PARAMS\n",
    "args['T'] = 25 # num steps to use (can use 10 or 15 steps to test), 25 is generally good, but use 50 or 100 for harder problems and better AF2 results\n",
    "\n",
    "# PARTIAL DIFFUSION\n",
    "args['sampling_temp'] = 1.0 # set to 1.0 for full diffusion, else entering partial diffusion mode, diffusing to specified temperature of input\n",
    "args['trb'] = None # path to trb if you want to partially diffuse and use the same contigs, else you can input a pdb or sequence to partially diffuse from\n",
    "\n",
    "# send args to sampler and initialize diffuser and model\n",
    "S.set_args(args)\n",
    "S.model_init()\n",
    "S.diffuser_init()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44f08fee-39fe-48ec-902c-5846aafc7b59",
   "metadata": {},
   "source": [
    "### Generate Samples\n",
    "Once you generate samples you like, use the args.json file saved to launch a production run\n",
    "\n",
    "For example:\n",
    "```\n",
    "python ./inference.py -input_json ./examples/out/design_000000_args.json\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5b625f99-ede1-4f6d-b2f0-f9829b921d13",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing sequence input\n",
      "Generating sample 000000 ...\n",
      "DAEELLGAEDALQLPKPEPSEQLAELREAAAKLKLRLA\n",
      "    TIMESTEP [25/25]   |   current PLDDT: 0.7031   <<  >>   best PLDDT: 0.7031\n",
      "SLRERLELLLELLRLLPEPSEQGIAELEELLEELLLLL\n",
      "    TIMESTEP [24/25]   |   current PLDDT: 0.8086   <<  >>   best PLDDT: 0.8086\n",
      "EEEEVKELEEIADKLPPEPSEQEREILEALEEQLLKIP\n",
      "    TIMESTEP [23/25]   |   current PLDDT: 0.8008   <<  >>   best PLDDT: 0.8086\n",
      "ELRDKKLQIAELFDALPEPSEQDLQREEQSLKKLLEKA\n",
      "    TIMESTEP [22/25]   |   current PLDDT: 0.8115   <<  >>   best PLDDT: 0.8115\n",
      "RRHLPEDLIRGLDRAVPEPSEQELIREREALVRELLGE\n",
      "    TIMESTEP [21/25]   |   current PLDDT: 0.8062   <<  >>   best PLDDT: 0.8115\n",
      "SALATDVLILLLAGELPEPSEQEALRLRERQEREAEAN\n",
      "    TIMESTEP [20/25]   |   current PLDDT: 0.8057   <<  >>   best PLDDT: 0.8115\n",
      "SEPQIAQPWEALLTVVPEPSEQRADTVRKQQLNALERF\n",
      "    TIMESTEP [19/25]   |   current PLDDT: 0.8062   <<  >>   best PLDDT: 0.8115\n",
      "SIKPLLEPSRRLLEGLPEPSEQQAAVEKLVELALLELR\n",
      "    TIMESTEP [18/25]   |   current PLDDT: 0.8545   <<  >>   best PLDDT: 0.8545\n",
      "ASGSALEEAKLKLRGIPEPSEQELELAERLERRLERRH\n",
      "    TIMESTEP [17/25]   |   current PLDDT: 0.8569   <<  >>   best PLDDT: 0.8569\n",
      "APKKKLAKKKILLQGIPEPSEQEIEEMLAEQEREKLIR\n",
      "    TIMESTEP [16/25]   |   current PLDDT: 0.8901   <<  >>   best PLDDT: 0.8901\n",
      "SPTKVLVRELAKAQGLPEPSEQELEQARALLEREQAEA\n",
      "    TIMESTEP [15/25]   |   current PLDDT: 0.8691   <<  >>   best PLDDT: 0.8901\n",
      "SPTKQALRAELKEAGLPEPSEQYLKQLKDLLEQELAIA\n",
      "    TIMESTEP [14/25]   |   current PLDDT: 0.8892   <<  >>   best PLDDT: 0.8901\n",
      "SPTLARLEEELKEAGLPEPSEQYLEELRDKAERELELS\n",
      "    TIMESTEP [13/25]   |   current PLDDT: 0.8682   <<  >>   best PLDDT: 0.8901\n",
      "KPSLERIEQELKEEGLPEPSEQYLKEIIEKAERLLALG\n",
      "    TIMESTEP [12/25]   |   current PLDDT: 0.9082   <<  >>   best PLDDT: 0.9082\n",
      "KPTLEEIRAELKAEGLPEPSEQQIKEAIEKKERLKELD\n",
      "    TIMESTEP [11/25]   |   current PLDDT: 0.9131   <<  >>   best PLDDT: 0.9131\n",
      "KLTLEEIREELKAEGLPEPSEQELREAIETIERRLELG\n",
      "    TIMESTEP [10/25]   |   current PLDDT: 0.9517   <<  >>   best PLDDT: 0.9517\n",
      "KLTLEEIREELKAEGLPEPSEQELREAIKTIERRLEQG\n",
      "    TIMESTEP [09/25]   |   current PLDDT: 0.9639   <<  >>   best PLDDT: 0.9639\n",
      "SLTLEEIKEELKANGLPEPSEQELRDALETIEREKEQG\n",
      "    TIMESTEP [08/25]   |   current PLDDT: 0.9238   <<  >>   best PLDDT: 0.9639\n",
      "KLTLEEIREELKANGWPEPSEQELRDALEKIEREREQG\n",
      "    TIMESTEP [07/25]   |   current PLDDT: 0.9678   <<  >>   best PLDDT: 0.9678\n",
      "ALTLDEIKEELKKNGWPEPSEQELRDALEKIERKREAG\n",
      "    TIMESTEP [06/25]   |   current PLDDT: 0.9683   <<  >>   best PLDDT: 0.9683\n",
      "ALTLDEIREELKKNGWPEPSEQELRDALEKIEREREQG\n",
      "    TIMESTEP [05/25]   |   current PLDDT: 0.9614   <<  >>   best PLDDT: 0.9683\n",
      "KLTLDEIREELKKNGWPEPSEQELRDALEKIDREREQG\n",
      "    TIMESTEP [04/25]   |   current PLDDT: 0.9668   <<  >>   best PLDDT: 0.9683\n",
      "KLTLDEIREELKKNGWPEPSEQELRDALEKIEREREQG\n",
      "    TIMESTEP [03/25]   |   current PLDDT: 0.9634   <<  >>   best PLDDT: 0.9683\n",
      "KLTLDEIREELKKNGWPEPSEQELRDALEKIEREREQG\n",
      "    TIMESTEP [02/25]   |   current PLDDT: 0.9399   <<  >>   best PLDDT: 0.9683\n",
      "KLTLDEIREELKKNGWPEPSEQELRDALEKIEREREQG\n",
      "    TIMESTEP [01/25]   |   current PLDDT: 0.9590   <<  >>   best PLDDT: 0.9683\n",
      "Finished design ./examples/out/design_000000 in 0.83 minutes.\n"
     ]
    }
   ],
   "source": [
    "for i in range(1):\n",
    "    S.generate_sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c2647e1-4019-4a36-bee1-708d3dbba247",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SE3nv",
   "language": "python",
   "name": "se3nv"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}