ProteinMPNN/colab_notebooks/quickdemo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "quickdemo.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dauparas/ProteinMPNN/blob/main/colab_notebooks/quickdemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#ProteinMPNN\n",
        "This notebook is intended as a quick demo, more features to come!"
      ],
      "metadata": {
        "id": "AYZebfKn8gef"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Setup Model\n",
        "import json, time, os, sys, glob\n",
        "\n",
        "if not os.path.isdir(\"ProteinMPNN\"):\n",
        "  os.system(\"git clone -q https://github.com/dauparas/ProteinMPNN.git\")\n",
        "sys.path.append('/content/ProteinMPNN/vanilla_proteinmpnn')\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import shutil\n",
        "import warnings\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch import optim\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.data.dataset import random_split, Subset\n",
        "import copy\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import random\n",
        "import os.path\n",
        "from protein_mpnn_utils import loss_nll, loss_smoothed, gather_edges, gather_nodes, gather_nodes_t, cat_neighbors_nodes, _scores, _S_to_seq, tied_featurize, parse_PDB\n",
        "from protein_mpnn_utils import StructureDataset, StructureDatasetPDB, ProteinMPNN\n",
        "\n",
        "device = torch.device(\"cuda:0\" if (torch.cuda.is_available()) else \"cpu\")\n",
        "model_name=\"v_48_020\"             # ProteinMPNN model name: v_48_002, v_48_010, v_48_020, v_48_030, v_32_002, v_32_010; v_32_020, v_32_030; v_48_010=version with 48 edges 0.10A noise\n",
        "backbone_noise=0.00               # Standard deviation of Gaussian noise to add to backbone atoms\n",
        "\n",
        "path_to_model_weights='/content/ProteinMPNN/vanilla_proteinmpnn/vanilla_model_weights'          \n",
        "hidden_dim = 128\n",
        "num_layers = 3 \n",
        "model_folder_path = path_to_model_weights\n",
        "if model_folder_path[-1] != '/':\n",
        "    model_folder_path = model_folder_path + '/'\n",
        "checkpoint_path = model_folder_path + f'{model_name}.pt'\n",
        "\n",
        "checkpoint = torch.load(checkpoint_path, map_location=device) \n",
        "print('Number of edges:', checkpoint['num_edges'])\n",
        "noise_level_print = checkpoint['noise_level']\n",
        "print(f'Training noise level: {noise_level_print}A')\n",
        "model = ProteinMPNN(num_letters=21, node_features=hidden_dim, edge_features=hidden_dim, hidden_dim=hidden_dim, num_encoder_layers=num_layers, num_decoder_layers=num_layers, augment_eps=backbone_noise, k_neighbors=checkpoint['num_edges'])\n",
        "model.to(device)\n",
        "model.load_state_dict(checkpoint['model_state_dict'])\n",
        "model.eval()\n",
        "print(\"Model loaded\")"
      ],
      "metadata": {
        "id": "iYDU3ftml2k5",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "from google.colab import files\n",
        "import numpy as np\n",
        "\n",
        "#########################\n",
        "def get_pdb(pdb_code=\"\"):\n",
        "  if pdb_code is None or pdb_code == \"\":\n",
        "    upload_dict = files.upload()\n",
        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
        "    return \"tmp.pdb\"\n",
        "  else:\n",
        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
        "    return f\"{pdb_code}.pdb\"\n",
        "\n",
        "#@markdown ### Input Options\n",
        "pdb='6MRR' #@param {type:\"string\"}\n",
        "pdb_path = get_pdb(pdb)\n",
        "#@markdown - pdb code (leave blank to get an upload prompt)\n",
        "\n",
        "designed_chain = \"A\" #@param {type:\"string\"}\n",
        "fixed_chain = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if designed_chain == \"\":\n",
        "  designed_chain_list = []\n",
        "else:\n",
        "  designed_chain_list = re.sub(\"[^A-Za-z]+\",\",\", designed_chain).split(\",\")\n",
        "\n",
        "if fixed_chain == \"\":\n",
        "  fixed_chain_list = []\n",
        "else:\n",
        "  fixed_chain_list = re.sub(\"[^A-Za-z]+\",\",\", fixed_chain).split(\",\")\n",
        "\n",
        "chain_list = list(set(designed_chain_list + fixed_chain_list))\n",
        "\n",
        "#@markdown - specified which chain(s) to design and which chain(s) to keep fixed. \n",
        "#@markdown   Use comma:`A,B` to specifiy more than one chain\n",
        "\n",
        "#chain = \"A\" #@param {type:\"string\"}\n",
        "#pdb_path_chains = chain\n",
        "##@markdown - Define which chain to redesign\n",
        "\n",
        "#@markdown ### Design Options\n",
        "num_seqs = 1 #@param [\"1\", \"2\", \"4\", \"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n",
        "num_seq_per_target = num_seqs\n",
        "sampling_temp = \"0.1\" #@param [\"0.1\", \"0.15\", \"0.2\", \"0.25\", \"0.3\"]\n",
        "#@markdown - Sampling temperature for amino acids, T=0.0 means taking \n",
        "#@markdown   argmax, T>>1.0 means sample randomly. Suggested values \n",
        "#@markdown   0.1, 0.15, 0.2, 0.25, 0.3. Higher values will lead to more diversity.\n",
        "\n",
        "\n",
        "save_score=0                      # 0 for False, 1 for True; save score=-log_prob to npy files\n",
        "save_probs=0                      # 0 for False, 1 for True; save MPNN predicted probabilites per position\n",
        "score_only=0                      # 0 for False, 1 for True; score input backbone-sequence pairs\n",
        "conditional_probs_only=0          # 0 for False, 1 for True; output conditional probabilities p(s_i given the rest of the sequence and backbone)\n",
        "conditional_probs_only_backbone=0 # 0 for False, 1 for True; if true output conditional probabilities p(s_i given backbone)\n",
        "    \n",
        "batch_size=1                      # Batch size; can set higher for titan, quadro GPUs, reduce this if running out of GPU memory\n",
        "max_length=20000                  # Max sequence length\n",
        "    \n",
        "out_folder='.'                    # Path to a folder to output sequences, e.g. /home/out/\n",
        "jsonl_path=''                     # Path to a folder with parsed pdb into jsonl\n",
        "omit_AAs='X'                      # Specify which amino acids should be omitted in the generated sequence, e.g. 'AC' would omit alanine and cystine.\n",
        "   \n",
        "pssm_multi=0.0                    # A value between [0.0, 1.0], 0.0 means do not use pssm, 1.0 ignore MPNN predictions\n",
        "pssm_threshold=0.0                # A value between -inf + inf to restric per position AAs\n",
        "pssm_log_odds_flag=0               # 0 for False, 1 for True\n",
        "pssm_bias_flag=0                   # 0 for False, 1 for True\n",
        "\n",
        "\n",
        "##############################################################\n",
        "\n",
        "folder_for_outputs = out_folder\n",
        "\n",
        "NUM_BATCHES = num_seq_per_target//batch_size\n",
        "BATCH_COPIES = batch_size\n",
        "temperatures = [float(item) for item in sampling_temp.split()]\n",
        "omit_AAs_list = omit_AAs\n",
        "alphabet = 'ACDEFGHIKLMNPQRSTVWYX'\n",
        "\n",
        "omit_AAs_np = np.array([AA in omit_AAs_list for AA in alphabet]).astype(np.float32)\n",
        "\n",
        "chain_id_dict = None\n",
        "fixed_positions_dict = None\n",
        "pssm_dict = None\n",
        "omit_AA_dict = None\n",
        "bias_AA_dict = None\n",
        "tied_positions_dict = None\n",
        "bias_by_res_dict = None\n",
        "bias_AAs_np = np.zeros(len(alphabet))\n",
        "\n",
        "\n",
        "###############################################################\n",
        "pdb_dict_list = parse_PDB(pdb_path, input_chain_list=chain_list)\n",
        "dataset_valid = StructureDatasetPDB(pdb_dict_list, truncate=None, max_length=max_length)\n",
        "\n",
        "chain_id_dict = {}\n",
        "chain_id_dict[pdb_dict_list[0]['name']]= (designed_chain_list, fixed_chain_list)\n",
        "\n",
        "print(chain_id_dict)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "k4o6w2Y23wxO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title RUN\n",
        "with torch.no_grad():\n",
        "  print('Generating sequences...')\n",
        "  for ix, protein in enumerate(dataset_valid):\n",
        "    score_list = []\n",
        "    all_probs_list = []\n",
        "    all_log_probs_list = []\n",
        "    S_sample_list = []\n",
        "    batch_clones = [copy.deepcopy(protein) for i in range(BATCH_COPIES)]\n",
        "    X, S, mask, lengths, chain_M, chain_encoding_all, chain_list_list, visible_list_list, masked_list_list, masked_chain_length_list_list, chain_M_pos, omit_AA_mask, residue_idx, dihedral_mask, tied_pos_list_of_lists_list, pssm_coef, pssm_bias, pssm_log_odds_all, bias_by_res_all, tied_beta = tied_featurize(batch_clones, device, chain_id_dict, fixed_positions_dict, omit_AA_dict, tied_positions_dict, pssm_dict, bias_by_res_dict)\n",
        "    pssm_log_odds_mask = (pssm_log_odds_all > pssm_threshold).float() #1.0 for true, 0.0 for false\n",
        "    name_ = batch_clones[0]['name']\n",
        "\n",
        "    randn_1 = torch.randn(chain_M.shape, device=X.device)\n",
        "    log_probs = model(X, S, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_1)\n",
        "    mask_for_loss = mask*chain_M*chain_M_pos\n",
        "    scores = _scores(S, log_probs, mask_for_loss)\n",
        "    native_score = scores.cpu().data.numpy()\n",
        "\n",
        "    for temp in temperatures:\n",
        "        for j in range(NUM_BATCHES):\n",
        "            randn_2 = torch.randn(chain_M.shape, device=X.device)\n",
        "            if tied_positions_dict == None:\n",
        "                sample_dict = model.sample(X, randn_2, S, chain_M, chain_encoding_all, residue_idx, mask=mask, temperature=temp, omit_AAs_np=omit_AAs_np, bias_AAs_np=bias_AAs_np, chain_M_pos=chain_M_pos, omit_AA_mask=omit_AA_mask, pssm_coef=pssm_coef, pssm_bias=pssm_bias, pssm_multi=pssm_multi, pssm_log_odds_flag=bool(pssm_log_odds_flag), pssm_log_odds_mask=pssm_log_odds_mask, pssm_bias_flag=bool(pssm_bias_flag), bias_by_res=bias_by_res_all)\n",
        "                S_sample = sample_dict[\"S\"] \n",
        "            else:\n",
        "                sample_dict = model.tied_sample(X, randn_2, S, chain_M, chain_encoding_all, residue_idx, mask=mask, temperature=temp, omit_AAs_np=omit_AAs_np, bias_AAs_np=bias_AAs_np, chain_M_pos=chain_M_pos, omit_AA_mask=omit_AA_mask, pssm_coef=pssm_coef, pssm_bias=pssm_bias, pssm_multi=pssm_multi, pssm_log_odds_flag=bool(pssm_log_odds_flag), pssm_log_odds_mask=pssm_log_odds_mask, pssm_bias_flag=bool(pssm_bias_flag), tied_pos=tied_pos_list_of_lists_list[0], tied_beta=tied_beta, bias_by_res=bias_by_res_all)\n",
        "            # Compute scores\n",
        "                S_sample = sample_dict[\"S\"]\n",
        "            log_probs = model(X, S_sample, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_2, use_input_decoding_order=True, decoding_order=sample_dict[\"decoding_order\"])\n",
        "            mask_for_loss = mask*chain_M*chain_M_pos\n",
        "            scores = _scores(S_sample, log_probs, mask_for_loss)\n",
        "            scores = scores.cpu().data.numpy()\n",
        "            all_probs_list.append(sample_dict[\"probs\"].cpu().data.numpy())\n",
        "            all_log_probs_list.append(log_probs.cpu().data.numpy())\n",
        "            S_sample_list.append(S_sample.cpu().data.numpy())\n",
        "            for b_ix in range(BATCH_COPIES):\n",
        "                masked_chain_length_list = masked_chain_length_list_list[b_ix]\n",
        "                masked_list = masked_list_list[b_ix]\n",
        "                seq_recovery_rate = torch.sum(torch.sum(torch.nn.functional.one_hot(S[b_ix], 21)*torch.nn.functional.one_hot(S_sample[b_ix], 21),axis=-1)*mask_for_loss[b_ix])/torch.sum(mask_for_loss[b_ix])\n",
        "                seq = _S_to_seq(S_sample[b_ix], chain_M[b_ix])\n",
        "                score = scores[b_ix]\n",
        "                score_list.append(score)\n",
        "                native_seq = _S_to_seq(S[b_ix], chain_M[b_ix])\n",
        "                if b_ix == 0 and j==0 and temp==temperatures[0]:\n",
        "                    start = 0\n",
        "                    end = 0\n",
        "                    list_of_AAs = []\n",
        "                    for mask_l in masked_chain_length_list:\n",
        "                        end += mask_l\n",
        "                        list_of_AAs.append(native_seq[start:end])\n",
        "                        start = end\n",
        "                    native_seq = \"\".join(list(np.array(list_of_AAs)[np.argsort(masked_list)]))\n",
        "                    l0 = 0\n",
        "                    for mc_length in list(np.array(masked_chain_length_list)[np.argsort(masked_list)])[:-1]:\n",
        "                        l0 += mc_length\n",
        "                        native_seq = native_seq[:l0] + '/' + native_seq[l0:]\n",
        "                        l0 += 1\n",
        "                    sorted_masked_chain_letters = np.argsort(masked_list_list[0])\n",
        "                    print_masked_chains = [masked_list_list[0][i] for i in sorted_masked_chain_letters]\n",
        "                    sorted_visible_chain_letters = np.argsort(visible_list_list[0])\n",
        "                    print_visible_chains = [visible_list_list[0][i] for i in sorted_visible_chain_letters]\n",
        "                    native_score_print = np.format_float_positional(np.float32(native_score.mean()), unique=False, precision=4)\n",
        "                    line = '>{}, score={}, fixed_chains={}, designed_chains={}, model_name={}\\n{}\\n'.format(name_, native_score_print, print_visible_chains, print_masked_chains, model_name, native_seq)\n",
        "                    print(line.rstrip())\n",
        "                start = 0\n",
        "                end = 0\n",
        "                list_of_AAs = []\n",
        "                for mask_l in masked_chain_length_list:\n",
        "                    end += mask_l\n",
        "                    list_of_AAs.append(seq[start:end])\n",
        "                    start = end\n",
        "\n",
        "                seq = \"\".join(list(np.array(list_of_AAs)[np.argsort(masked_list)]))\n",
        "                l0 = 0\n",
        "                for mc_length in list(np.array(masked_chain_length_list)[np.argsort(masked_list)])[:-1]:\n",
        "                    l0 += mc_length\n",
        "                    seq = seq[:l0] + '/' + seq[l0:]\n",
        "                    l0 += 1\n",
        "                score_print = np.format_float_positional(np.float32(score), unique=False, precision=4)\n",
        "                seq_rec_print = np.format_float_positional(np.float32(seq_recovery_rate.detach().cpu().numpy()), unique=False, precision=4)\n",
        "                line = '>T={}, sample={}, score={}, seq_recovery={}\\n{}\\n'.format(temp,b_ix,score_print,seq_rec_print,seq)\n",
        "                print(line.rstrip())\n",
        "\n",
        "\n",
        "all_probs_concat = np.concatenate(all_probs_list)\n",
        "all_log_probs_concat = np.concatenate(all_log_probs_list)\n",
        "S_sample_concat = np.concatenate(S_sample_list)"
      ],
      "metadata": {
        "id": "xMVlYh8Fv2of",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# experimental output\n",
        "plt.figure(figsize=(20,5), dpi=100)\n",
        "plt.imshow(all_probs_concat.mean(0).T,vmin=0,vmax=1)\n",
        "plt.xlabel(\"positions\")\n",
        "plt.ylabel(\"amino acids\")\n",
        "plt.yticks(range(21),list(alphabet))\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "4jSKLU3L17Sf"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}