2023_12_11_rla_utils.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0fb47531-1fa8-4d32-9b71-b8e104022fab",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-07 16:02:22.564563: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-03-07 16:02:22.712068: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-03-07 16:02:23.779901: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory\n",
      "2024-03-07 16:02:23.779966: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory\n",
      "2024-03-07 16:02:23.779972: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "/home/gridsan/fbirnbaum/.local/lib/python3.9/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n",
      "/state/partition1/slurm_tmp/25344618.0.0/ipykernel_238629/2500981057.py:21: DeprecationWarning: Please use `pearsonr` from the `scipy.stats` namespace, the `scipy.stats.stats` namespace is deprecated.\n",
      "  from scipy.stats.stats import pearsonr\n"
     ]
    }
   ],
   "source": [
    "import src.models_and_optimizers as model_utils\n",
    "import yaml\n",
    "from types import SimpleNamespace\n",
    "from clip_main import get_wds_loaders\n",
    "from transformers import EsmTokenizer\n",
    "import src.data_utils as data_utils\n",
    "import os\n",
    "import torch\n",
    "import sys\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "from tmtools.io import get_structure, get_residue_data\n",
    "from tmtools import tm_align\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "from torch.cuda.amp import autocast\n",
    "import tmscoring\n",
    "import json\n",
    "import copy\n",
    "from scipy.stats.stats import pearsonr \n",
    "import matplotlib.pyplot as plt\n",
    "import yaml\n",
    "import pandas as pd\n",
    "import glob\n",
    "import webdataset as wds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5c422249-f113-44d2-b783-b93842f270fb",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'torch' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 207\u001b[0m\n\u001b[1;32m    199\u001b[0m     gnn_features, text_features, logit_scale \u001b[38;5;241m=\u001b[39m model(text_inp, coord_data)\n\u001b[1;32m    200\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m {\n\u001b[1;32m    201\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtext\u001b[39m\u001b[38;5;124m'\u001b[39m: text_features, \u001b[38;5;66;03m# text feature\u001b[39;00m\n\u001b[1;32m    202\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mgnn\u001b[39m\u001b[38;5;124m'\u001b[39m: gnn_features, \u001b[38;5;66;03m# gnn feature\u001b[39;00m\n\u001b[1;32m    203\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mseq_mask_with_burn_in\u001b[39m\u001b[38;5;124m'\u001b[39m: seq_batch[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mseq_loss_mask\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;241m0\u001b[39m], \u001b[38;5;66;03m# sequence mask of what's supervised\u001b[39;00m\n\u001b[1;32m    204\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcoord_mask_with_burn_in\u001b[39m\u001b[38;5;124m'\u001b[39m: coords_batch[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcoords_loss_mask\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;241m0\u001b[39m], \u001b[38;5;66;03m# coord mask of what's supervised\u001b[39;00m\n\u001b[1;32m    205\u001b[0m     }\n\u001b[0;32m--> 207\u001b[0m cos \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mCosineSimilarity()\n\u001b[1;32m    208\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcalc_sim\u001b[39m(all_outputs):\n\u001b[1;32m    209\u001b[0m     all_sims \u001b[38;5;241m=\u001b[39m []\n",
      "\u001b[0;31mNameError\u001b[0m: name 'torch' is not defined"
     ]
    }
   ],
   "source": [
    "def mask_peptide(seq_batch, coords_batch, pdb):\n",
    "    chain_len_dicts = {}\n",
    "    chains, lens = torch.unique_consecutive(coords_batch['chain_lens'][0][0], return_counts=True)\n",
    "    chain_len_dicts['protein'] = torch.max(lens)\n",
    "    chain_len_dicts['peptide'] = torch.min(lens)\n",
    "    peptide_len = chain_len_dicts['peptide']\n",
    "    from_back = torch.argmin(lens) == 1\n",
    "    if from_back:\n",
    "        seq_batch['string_sequence'][0] = seq_batch['string_sequence'][0][:-peptide_len] + 'X'*peptide_len\n",
    "        seq_batch['seq_loss_mask'][0][:,:-1*peptide_len] = False\n",
    "        seq_batch['seq_loss_mask'][1][:,:-1*peptide_len] = False\n",
    "    else:\n",
    "        seq_batch['string_sequence'][0] = 'X'*peptide_len + seq_batch['string_sequence'][0][peptide_len:]\n",
    "        seq_batch['seq_loss_mask'][0][:,peptide_len:] = False\n",
    "        seq_batch['seq_loss_mask'][1][:,peptide_len:] = False\n",
    "    return seq_batch\n",
    "\n",
    "def mask_all(seq_batch, pdb):\n",
    "    seq_batch['string_sequence'][0] = 'X'*len(seq_batch['string_sequence'][0])\n",
    "    seq_batch['seq_loss_mask'][0][:,:] = False\n",
    "    seq_batch['seq_loss_mask'][1][:,:] = False\n",
    "    return seq_batch\n",
    "\n",
    "# Extract kNN info\n",
    "def extract_knn(X, eps, top_k):\n",
    "    # Convolutional network on NCHW\n",
    "    dX = torch.unsqueeze(X, 1) - torch.unsqueeze(X, 2)\n",
    "    D = torch.sqrt(torch.sum(dX**2, 3) + eps)\n",
    "    mask_2D = torch.ones(D.shape)\n",
    "    D *= mask_2D\n",
    "\n",
    "    # Identify k nearest neighbors (including self)\n",
    "    D_max, _ = torch.max(D, -1, keepdim=True)\n",
    "    D_adjust = D + (1. - mask_2D) * D_max\n",
    "    D_neighbors, E_idx = torch.topk(D_adjust, top_k, dim=-1, largest=False)\n",
    "    return D_neighbors, E_idx\n",
    "\n",
    "def get_interaction_res(coords_batch, pdb, top_k, threshold=5, remove_far=False, prot_only=False):\n",
    "    chain_len_dicts = {}\n",
    "    chains, lens = torch.unique_consecutive(coords_batch['chain_lens'][0][0], return_counts=True)\n",
    "    chain_len_dicts['protein'] = torch.max(lens)\n",
    "    chain_len_dicts['peptide'] = torch.min(lens)\n",
    "    from_back = torch.argmin(lens) == 1\n",
    "    top_k = min(top_k, coords_batch['coords'][0].size(1))\n",
    "    D_neighbors, E_idx = extract_knn(coords_batch['coords'][0][:,:,1,:], eps=1e-6, top_k=top_k) \n",
    "    if from_back:\n",
    "        interaction_res = set(range(chain_len_dicts['protein'], chain_len_dicts['protein']+chain_len_dicts['peptide']))\n",
    "    else:\n",
    "        interaction_res = set(range(0, chain_len_dicts['peptide']))\n",
    "    # interaction_res = set(range(0, chain_len_dicts['protein']+chain_len_dicts['peptide']))\n",
    "    if top_k == 0:\n",
    "        return list(interaction_res)\n",
    "    prot_to_add = set()\n",
    "    for res in range(0, chain_len_dicts['protein']+chain_len_dicts['peptide']):\n",
    "        if res in interaction_res:\n",
    "            prot_to_add = prot_to_add.union(set(E_idx[0, res].tolist()))\n",
    "        else:\n",
    "            for r in E_idx[0, res]:\n",
    "                if r in interaction_res:\n",
    "                    prot_to_add.add(r)\n",
    "    interaction_res = list(interaction_res.union(prot_to_add))\n",
    "    if remove_far:\n",
    "        to_remove = []\n",
    "        for res in interaction_res:\n",
    "            nother = 0\n",
    "            opp = 1 - coords_batch['chain_lens'][0][0][res].item()\n",
    "            for nres in E_idx[0, res]:\n",
    "                if coords_batch['chain_lens'][0][0][nres].item() == opp:\n",
    "                    nother += 1\n",
    "\n",
    "            if nother < threshold:\n",
    "                to_remove.append(res)\n",
    "        for res in to_remove:\n",
    "            interaction_res.remove(res)\n",
    "            \n",
    "    # Remove masked residues\n",
    "    to_remove = []\n",
    "    for res in interaction_res:\n",
    "        if not coords_batch['coords'][1][0, res]:\n",
    "            to_remove.append(res)\n",
    "    for res in to_remove:\n",
    "        interaction_res.remove(res)\n",
    "        \n",
    "    # Remove peptide residues if required\n",
    "    if prot_only:\n",
    "        to_remove = []\n",
    "        for res in interaction_res:\n",
    "            if (from_back and res in list(range(chain_len_dicts['protein'], chain_len_dicts['protein']+chain_len_dicts['peptide']))) or (not from_back and res in list(range(0, chain_len_dicts['peptide']))):\n",
    "                to_remove.append(res)\n",
    "        for res in to_remove:\n",
    "            interaction_res.remove(res)\n",
    "    return interaction_res\n",
    "\n",
    "def get_inter_dists(coords_batch, interaction_res, eps=1e-6, norm=True):\n",
    "    chain_len_dicts = {}\n",
    "    chains, lens = torch.unique_consecutive(coords_batch['chain_lens'][0][0], return_counts=True)\n",
    "    chain_len_dicts['protein'] = torch.max(lens)\n",
    "    chain_len_dicts['peptide'] = torch.min(lens)\n",
    "    from_back = torch.argmin(lens) == 1\n",
    "    if from_back:\n",
    "        pep_res = list(range(chain_len_dicts['protein'], chain_len_dicts['protein']+chain_len_dicts['peptide']))\n",
    "        prot_res = list(range(0, chain_len_dicts['protein']))\n",
    "    else:\n",
    "        pep_res = list(range(0, chain_len_dicts['peptide']))\n",
    "        prot_res = list(range(chain_len_dicts['peptide'], chain_len_dicts['peptide']+chain_len_dicts['protein']))\n",
    "    pep_coords = coords_batch['coords'][0][:,pep_res,1,:]\n",
    "    prot_coords = coords_batch['coords'][0][:,prot_res,1,:]\n",
    "    inter_dists = []\n",
    "    for res in range(chain_len_dicts['protein'] + chain_len_dicts['peptide']):\n",
    "        if not res in interaction_res:\n",
    "            inter_dists.append(np.nan)\n",
    "            continue\n",
    "        res_coord = coords_batch['coords'][0][:,res,1,:]\n",
    "        if res in pep_res:\n",
    "            dists = torch.sqrt(torch.sum((res_coord - prot_coords)**2, dim=2))\n",
    "        else:\n",
    "            dists = torch.sqrt(torch.sum((res_coord - pep_coords)**2, dim=2))\n",
    "        inter_dists.append(torch.min(dists))\n",
    "    inter_dists = np.array(inter_dists)\n",
    "    # inter_dists = (np.nanmax(inter_dists) - inter_dists)\n",
    "    if norm:\n",
    "        inter_dists = inter_dists / np.nanmax(inter_dists)\n",
    "    # inter_dists -= np.nanmin(inter_dists)\n",
    "    # inter_dists = 1 - (inter_dists / np.nanmax(inter_dists))\n",
    "    # inter_dists -= np.nanmean(inter_dists)\n",
    "    # inter_dists = np.abs(inter_dists) / np.nanstd(np.abs(inter_dists))\n",
    "    inter_dists = np.nan_to_num(inter_dists, nan=0)\n",
    "    return inter_dists\n",
    "            \n",
    "\n",
    "def mask_peptide_struct(coord_data, coords_batch, pdb, mask_prot=False):\n",
    "    peptide_res = get_interaction_res(coords_batch, pdb, top_k=0)\n",
    "    if not mask_prot:\n",
    "        mask = torch.ones(coord_data['x_mask'].shape).to(dtype=coord_data['x_mask'].dtype, device=coord_data['x_mask'].device)\n",
    "        mask[:,peptide_res] = 0\n",
    "    else:\n",
    "        mask = torch.zeros(coord_data['x_mask'].shape).to(dtype=coord_data['x_mask'].dtype, device=coord_data['x_mask'].device)\n",
    "        mask[:,peptide_res] = 1\n",
    "    coord_data['x_mask'] *= mask\n",
    "    # coord_data['X'] *= mask.unsqueeze(-1).unsqueeze(-1)\n",
    "    return coord_data\n",
    "    \n",
    "def get_text_and_image_features(model, tokenizer, batch, pdb=None, seq_mask=None, struct_mask=None, focus=None, top_k=30, remove_far=False, threshold=5, weight_dists=False, get_peptide_mask=True, dev='cuda:0'):\n",
    "    seq_batch, coords_batch = batch\n",
    "    if seq_mask == 'peptide':\n",
    "        seq_batch = mask_peptide(seq_batch, coords_batch, pdb)\n",
    "    if seq_mask == 'all':\n",
    "        seq_batch = mask_all(seq_batch, pdb)\n",
    "    seqs = seq_batch['string_sequence']\n",
    "    text_inp = tokenizer(seqs, return_tensors='pt', padding=True, truncation=True, max_length=1024+2)\n",
    "    text_inp['position_ids'] = seq_batch['pos_embs'][0]\n",
    "    text_inp = {k: v.to(dev) for k, v in text_inp.items()}\n",
    "    coord_data = data_utils.construct_gnn_inp(coords_batch, device=dev, half_precision=True)\n",
    "    if struct_mask=='peptide':\n",
    "        coord_data = mask_peptide_struct(coord_data, coords_batch, pdb)\n",
    "    elif struct_mask=='protein':\n",
    "        coord_data = mask_peptide_struct(coord_data, coords_batch, pdb, mask_prot=True)\n",
    "    gnn_features, text_features, logit_scale = model(text_inp, coord_data)\n",
    "    new_text_features, _, new_text_mask = data_utils.postprocess_text_features(\n",
    "        text_features=text_features, \n",
    "        inp_dict=text_inp, \n",
    "        tokenizer=tokenizer, \n",
    "        placeholder_mask=seq_batch['placeholder_mask'][0])\n",
    "    if focus:\n",
    "        interaction_res = get_interaction_res(coords_batch, pdb, top_k, remove_far = remove_far, threshold = threshold)\n",
    "        interaction_mask = torch.zeros(new_text_mask.shape).to(dtype=new_text_mask.dtype, device=new_text_mask.device)\n",
    "        interaction_mask[:,interaction_res] = 1\n",
    "        new_text_mask *= interaction_mask\n",
    "    if weight_dists:\n",
    "        weights = torch.from_numpy(get_inter_dists(coords_batch, interaction_res)).to(device=gnn_features.device)\n",
    "    else:\n",
    "        weights = None\n",
    "    if get_peptide_mask:\n",
    "        peptide_mask_ind = get_interaction_res(coords_batch, pdb, 0)\n",
    "        peptide_mask = torch.zeros_like(new_text_mask.bool())\n",
    "        peptide_mask[0][peptide_mask_ind] = True\n",
    "    else:\n",
    "        peptide_mask = None\n",
    "    return {\n",
    "        'text': new_text_features, # text feature\n",
    "        'gnn': gnn_features, # gnn feature\n",
    "        'seq_mask_with_burn_in': seq_batch['seq_loss_mask'][0], # sequence mask of what's supervised\n",
    "        'coord_mask_with_burn_in': coords_batch['coords_loss_mask'][0], # coord mask of what's supervised\n",
    "        'seq_mask_no_burn_in': new_text_mask.bool(), # sequence mask of what's valid (e.g., not padded)\n",
    "        'coord_mask_no_burn_in': coords_batch['coords'][1], # coord mask of what's valid\n",
    "        'weights': weights, # distance weights, if requested\n",
    "        'peptide_mask': peptide_mask,\n",
    "        'coord_data': coord_data\n",
    "    }\n",
    "\n",
    "def get_text_and_image_features_clip(model, tokenizer, batch, dev='cuda'):\n",
    "    seq_batch, coords_batch = batch\n",
    "    seqs = seq_batch['string_sequence']\n",
    "    text_inp = tokenizer(seqs, return_tensors='pt', padding=True, truncation=True, max_length=1024+2)\n",
    "    text_inp['position_ids'] = seq_batch['pos_embs'][0]\n",
    "    text_inp = {k: v.to(dev) for k, v in text_inp.items()}\n",
    "    coord_data = data_utils.construct_gnn_inp(coords_batch, device=dev, half_precision=True)\n",
    "    gnn_features, text_features, logit_scale = model(text_inp, coord_data)\n",
    "    return {\n",
    "        'text': text_features, # text feature\n",
    "        'gnn': gnn_features, # gnn feature\n",
    "        'seq_mask_with_burn_in': seq_batch['seq_loss_mask'][0], # sequence mask of what's supervised\n",
    "        'coord_mask_with_burn_in': coords_batch['coords_loss_mask'][0], # coord mask of what's supervised\n",
    "    }\n",
    "\n",
    "cos = torch.nn.CosineSimilarity()\n",
    "def calc_sim(all_outputs):\n",
    "    all_sims = []\n",
    "    all_sims_burn = []\n",
    "    for output in all_outputs:\n",
    "        t = output['text'][output['seq_mask_no_burn_in']]\n",
    "        g = output['gnn'][output['coord_mask_no_burn_in']]\n",
    "        sim = (t.unsqueeze(1) @ g.unsqueeze(-1)).squeeze(1).squeeze(1)\n",
    "        all_sims.append(torch.mean(sim))\n",
    "\n",
    "    return all_sims\n",
    "\n",
    "def calc_sim_clip(all_outputs):\n",
    "    all_sims = []\n",
    "    all_sims_burn = []\n",
    "    for output in all_outputs:\n",
    "        t = output['text']\n",
    "        g = output['gnn']\n",
    "        ### ??? ###\n",
    "        sim = (t.unsqueeze(1) @ g.unsqueeze(-1)).squeeze(1).squeeze(1)\n",
    "        all_sims.append(torch.mean(sim))\n",
    "        ### ??? ###\n",
    "    return all_sims\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3ad7aefe-2d89-4a54-b817-145bffc14c45",
   "metadata": {},
   "outputs": [],
   "source": [
    "## ESM mutation analysis functions\n",
    "def get_muts(wt, mut):\n",
    "    inds = []\n",
    "    muts = []\n",
    "    for i, (wchar, mchar) in enumerate(zip(wt, mut)):\n",
    "        if wchar != mchar:\n",
    "            inds.append(i)\n",
    "            muts.append(mchar)\n",
    "    return inds, muts\n",
    "\n",
    "def score_mut(wt, idx, mt, token_probs, alphabet):\n",
    "    wt = wt[idx]\n",
    "    wt_encoded, mt_encoded = alphabet.get_idx(wt), alphabet.get_idx(mt)\n",
    "\n",
    "    # add 1 for BOS\n",
    "    score = token_probs[0, 1 + idx, mt_encoded] - token_probs[0, 1 + idx, wt_encoded]\n",
    "    return score.item()\n",
    "\n",
    "def score_protein(idxs, mts, wt_seq, model, alphabet):\n",
    "\n",
    "    # inference for each model\n",
    "    batch_converter = alphabet.get_batch_converter()\n",
    "\n",
    "    data = [\n",
    "        (\"protein1\", wt_seq),\n",
    "    ]\n",
    "    batch_labels, batch_strs, batch_tokens = batch_converter(data)\n",
    "\n",
    "    all_token_probs = []\n",
    "    for i in tqdm(range(batch_tokens.size(1))):\n",
    "        batch_tokens_masked = batch_tokens.clone()\n",
    "        batch_tokens_masked[0, i] = alphabet.mask_idx\n",
    "        with torch.no_grad():\n",
    "            token_probs = torch.log_softmax(\n",
    "                model(batch_tokens_masked.cuda())[\"logits\"], dim=-1\n",
    "            )\n",
    "        all_token_probs.append(token_probs[:, i])  # vocab size\n",
    "    token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)\n",
    "    esm_predictions = []\n",
    "    for idx_list, mut_list in zip(idxs, mts):\n",
    "        mut_score = 0\n",
    "        for idx, mut in zip(idx_list, mut_list):\n",
    "            mut_score += score_mut(wt_seq, idx, mut, token_probs, alphabet)\n",
    "        if len(idx_list) == 0:\n",
    "            mut_score = np.nan\n",
    "        esm_predictions.append(mut_score)\n",
    "    return esm_predictions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18abb9a5-ad7f-4535-8f62-643e695bd5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_score(output_dict, weight_dists, MAX_LEN, pep_weight=1, plot_scores=None, plot_weights=None, plot_pep_mask=None, plot_indices=None, plot_X=None, plot_seq=None, is_complex=False):\n",
    "    text_feat = output_dict['text']\n",
    "    gnn_feat =  output_dict['gnn'][:, :text_feat.shape[1]] # remove tail padding\n",
    "    scores = (text_feat.unsqueeze(2) @ gnn_feat.unsqueeze(-1)).squeeze(-1).squeeze(-1)\n",
    "    if is_complex:\n",
    "        pep_scores = list(scores[output_dict['peptide_mask']].cpu().numpy())\n",
    "        prot_scores = list(scores[~output_dict['peptide_mask']].cpu().numpy())\n",
    "        # print('\\t', scores.shape, scores[output_dict['peptide_mask']].cpu().shape, scores[~output_dict['peptide_mask']].cpu().shape)\n",
    "        pep_len = len(pep_scores)\n",
    "        inter_len = (MAX_LEN - pep_len - len(prot_scores))\n",
    "        plot_pep_mask += list(output_dict['peptide_mask'][0][:pep_len].cpu().numpy()) + inter_len*[True] + list(output_dict['peptide_mask'][0][pep_len:].cpu().numpy())\n",
    "        pep_scores += inter_len*[-1]\n",
    "        plot_scores += (pep_scores + prot_scores)\n",
    "        plot_indices += list(range(MAX_LEN)) #output_dict['peptide_mask'].shape[1]\n",
    "    \n",
    "    X = batch[1]['coords'][0][0].cpu().numpy()\n",
    "    # print(inter_len, pep_len, len(prot_scores))\n",
    "    if is_complex and inter_len > 0:\n",
    "        plot_X.append(np.concatenate((X[:pep_len], -1*np.ones((inter_len, X.shape[1], X.shape[2])), X[pep_len:]), axis=0))\n",
    "    elif is_complex and inter_len == 0:\n",
    "        plot_X.append(X)\n",
    "    if is_complex and plot_X[-1].shape[0] != MAX_LEN:\n",
    "        print(plot_X[-1].shape[0])\n",
    "        print(inter_len)\n",
    "        raise ValueError\n",
    "    if is_complex:\n",
    "        plot_seq.append(batch[0]['string_sequence'][0][:pep_len] + 'X'*inter_len + batch[0]['string_sequence'][0][pep_len+25:])\n",
    "    if weight_dists:\n",
    "        # weights = output_dict['weights'] / torch.sum( output_dict['weights'])\n",
    "        if is_complex:\n",
    "            plot_weights += list(output_dict['weights'][:pep_len].cpu().numpy()) + (MAX_LEN - pep_len - len(prot_scores))*[-1] + list(output_dict['weights'][pep_len:].cpu().numpy())\n",
    "        if len(plot_scores) != len(plot_weights) or len(plot_indices) != len(plot_weights):\n",
    "            print(len(plot_scores), len(plot_weights), len(plot_indices))\n",
    "            raise ValueError\n",
    "        # output_dict['weights'][:] = 0\n",
    "        # output_dict['weights'][10] = 1\n",
    "        # cur_weights = torch.from_numpy(np.concatenate([var_weights[:pep_len], var_weights[len(pep_scores):]])).to(output_dict['weights'].device)\n",
    "        cur_weights = output_dict['weights']\n",
    "        # scores *= cur_weights\n",
    "        if is_complex:\n",
    "            cur_pep_weights = cur_weights * output_dict['peptide_mask'][0]\n",
    "            cur_pep_weights = cur_pep_weights / torch.sum(cur_pep_weights)\n",
    "            cur_prot_weights = cur_weights * ~output_dict['peptide_mask'][0]\n",
    "            cur_prot_weights = cur_prot_weights / torch.sum(cur_prot_weights)\n",
    "            cur_weights = cur_pep_weights + cur_prot_weights\n",
    "        scores *= cur_weights\n",
    "        # scores += scores*cur_weights # output_dict['weights']\n",
    "        # scores /= 2\n",
    "\n",
    "    if is_complex:\n",
    "        pep_scores = scores * output_dict['peptide_mask']\n",
    "        pep_seq_mask = output_dict['seq_mask_no_burn_in'].float() * output_dict['peptide_mask']\n",
    "        prot_scores = scores * ~output_dict['peptide_mask']\n",
    "        prot_seq_mask = output_dict['seq_mask_no_burn_in'].float() * ~output_dict['peptide_mask']\n",
    "        # special_pep_mask = torch.ones_like(output_dict['seq_mask_no_burn_in']).to(dtype=torch.float32)\n",
    "        # special_pep_mask[0,:18] = 1\n",
    "        # if design_set == design_sets[-1]:\n",
    "        #     pep_scores *= special_pep_mask\n",
    "        scores = pep_scores + prot_scores\n",
    "        pep_score = pep_weight*(pep_scores * pep_seq_mask).sum(1)/pep_seq_mask.sum(1)\n",
    "        prot_score = (prot_scores * prot_seq_mask).sum(1)/prot_seq_mask.sum(1)\n",
    "        if not torch.isnan(prot_score).item():\n",
    "            score = (pep_score.cpu().item() + prot_score.cpu().item()) / 2\n",
    "        else:\n",
    "            score = pep_score.cpu().item()\n",
    "    else:\n",
    "        score = (scores * output_dict['seq_mask_no_burn_in'].float()).sum(1)/output_dict['seq_mask_no_burn_in'].sum(1)\n",
    "        score = score.cpu().item()\n",
    "    return score, scores, plot_scores, plot_weights, plot_pep_mask, plot_indices, plot_X, plot_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "14b29938-c8b4-4723-86c5-58ed060f19cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def append_to_str(batch):\n",
    "    # batch = copy.deepcopy(batch)\n",
    "    batch[0]['string_sequence'] = [5*'G' + batch[0]['string_sequence'][0]]\n",
    "    batch_pad = torch.zeros(1, 5).to(dtype=batch[0]['pos_embs'][0].dtype, device=batch[0]['pos_embs'][0].device)\n",
    "    batch_pad_mask = torch.ones(1, 5).to(dtype=batch[0]['pos_embs'][1].dtype, device=batch[0]['pos_embs'][1].device)\n",
    "    batch[0]['pos_embs'] = [torch.cat([-1 + batch_pad, batch[0]['pos_embs'][0]], 1).to(dtype=batch[0]['pos_embs'][0].dtype),\n",
    "                                torch.cat([batch_pad_mask, batch[0]['pos_embs'][1]], 1).to(dtype=batch[0]['pos_embs'][1].dtype)]\n",
    "    batch[0]['placeholder_mask'] = [torch.cat([batch_pad_mask, batch[0]['placeholder_mask'][0]], 1).to(dtype=batch[0]['placeholder_mask'][0].dtype), \n",
    "                                        torch.cat([batch_pad_mask, batch[0]['placeholder_mask'][1]], 1).to(dtype=batch[0]['placeholder_mask'][1].dtype)]\n",
    "    batch[0]['seq_loss_mask'] = [torch.cat([batch_pad_mask, batch[0]['seq_loss_mask'][0]], 1).to(dtype=batch[0]['seq_loss_mask'][0].dtype),\n",
    "                                     torch.cat([batch_pad_mask, batch[0]['seq_loss_mask'][1]], 1).to(dtype=batch[0]['seq_loss_mask'][1].dtype)]\n",
    "    batch[0]['seq_to_coords'] = [torch.cat([-1 + batch_pad, batch[0]['seq_to_coords'][0]], 1).to(dtype=batch[0]['seq_to_coords'][0].dtype),\n",
    "                                     torch.cat([batch_pad_mask, batch[0]['seq_to_coords'][1]], 1).to(dtype=batch[0]['seq_to_coords'][0].dtype)]\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f36f73-c0a8-42d8-9805-fc29030ee0b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reverse_batch(batch):\n",
    "    batch = copy.deepcopy(batch)\n",
    "    batch[0]['string_sequence'] = [batch[0]['string_sequence'][0][::-1]]\n",
    "    # batch[0]['pos_embs'] = [torch.flip(batch[0]['pos_embs'][0], 1), torch.flip(batch[0]['pos_embs'][1], 1)]\n",
    "    # batch[0]['placeholder_mask'] = [torch.flip(batch[0]['placeholder_mask'][0], 1), torch.flip(batch[0]['placeholder_mask'][1], 1)]\n",
    "    # batch[0]['seq_loss_mask'] = [torch.flip(batch[0]['seq_loss_mask'][0], 1), torch.flip(batch[0]['seq_loss_mask'][1], 1)]\n",
    "    # batch[0]['seq_to_coords'] = [torch.flip(batch[0]['seq_to_coords'][0], 1), torch.flip(batch[0]['seq_to_coords'][1], 1)]\n",
    "    batch[1]['coords'] = [torch.flip(batch[1]['coords'][0], [1]).to(dtype=batch[1]['coords'][0].dtype, device=batch[1]['coords'][0].device), torch.flip(batch[1]['coords'][1], [1]).to(dtype=batch[1]['coords'][1].dtype, device=batch[1]['coords'][1].device)]\n",
    "    # batch[1]['res_info'] = [batch[1]['res_info'][0][::-1]]\n",
    "    # batch[1]['seq_lens'] = [batch[1]['seq_lens'][0][::-1]]\n",
    "    # batch[1]['coords_loss_mask'] = [torch.flip(batch[1]['coords_loss_mask'][0], 1), torch.flip(batch[1]['coords_loss_mask'][1], 1)]\n",
    "    # batch[1]['coords_to_seq'] = [torch.flip(batch[1]['coords_to_seq'][0], 1), torch.flip(batch[1]['coords_to_seq'][1], 1)]\n",
    "    # batch[1]['chain_dict']['begin'] = [batch[1]['chain_dict']['begin'][0].narrow(1, 0, pep_len), batch[1]['chain_dict']['begin'][1].narrow(1, 0, pep_len)]\n",
    "    # batch[1]['chain_dict']['end'] = [batch[1]['chain_dict']['end'][0].narrow(1, 0, pep_len), batch[1]['chain_dict']['end'][1].narrow(1, 0, pep_len)]\n",
    "    # batch[1]['chain_dict']['end'][0][0,-1] = batch[1]['chain_dict']['end'][0][0,-2] \n",
    "    # batch[1]['chain_dict']['singles'] = [batch[1]['chain_dict']['singles'][0].narrow(1, 0, pep_len), batch[1]['chain_dict']['singles'][1].narrow(1, 0, pep_len)]\n",
    "    # batch[1]['chain_dict']['ids'] = [batch[1]['chain_dict']['ids'][0].narrow(1, 0, pep_len) - pep_id, batch[1]['chain_dict']['ids'][1].narrow(1, 0, pep_len)]\n",
    "    # batch[1]['chain_lens'] = [batch[1]['chain_lens'][0].narrow(1, 0, pep_len) - pep_id, batch[1]['chain_lens'][1].narrow(1, 0, pep_len)]\n",
    "    return batch\n",
    "    \n",
    "def reverse_batch_chain(batch, pep_len):\n",
    "    batch = copy.deepcopy(batch)\n",
    "    batch[0]['string_sequence'] = [batch[0]['string_sequence'][0][:pep_len][::-1] + batch[0]['string_sequence'][0][pep_len:]]\n",
    "    batch[1]['coords'] = [torch.cat([torch.flip(batch[1]['coords'][0][:,:pep_len], [1]), batch[1]['coords'][0][:,pep_len:]], 1).to(dtype=batch[1]['coords'][0].dtype, device=batch[1]['coords'][0].device),\n",
    "                          torch.cat([torch.flip(batch[1]['coords'][1][:,:pep_len], [1]), batch[1]['coords'][1][:,pep_len:]], 1).to(dtype=batch[1]['coords'][1].dtype, device=batch[1]['coords'][1].device)]\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69e41f16-cfff-44d1-8d3f-d6f788cfc9e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def segment_batch(pep_batch, prot_batch, pep_id, prot_id, pep_len, prot_len):\n",
    "    pep_batch[0]['string_sequence'] = [pep_batch[0]['string_sequence'][0][:pep_len]]\n",
    "    pep_batch[0]['pos_embs'] = [pep_batch[0]['pos_embs'][0].narrow(1, 0, pep_len+2), pep_batch[0]['pos_embs'][1].narrow(1, 0, pep_len+2)]\n",
    "    pep_batch[0]['placeholder_mask'] = [pep_batch[0]['placeholder_mask'][0].narrow(1, 0, pep_len), pep_batch[0]['placeholder_mask'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[0]['seq_loss_mask'] = [pep_batch[0]['seq_loss_mask'][0].narrow(1, 0, pep_len), pep_batch[0]['seq_loss_mask'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[0]['seq_to_coords'] = [pep_batch[0]['seq_to_coords'][0].narrow(1, 0, pep_len), pep_batch[0]['seq_to_coords'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['coords'] = [pep_batch[1]['coords'][0].narrow(1, 0, pep_len), pep_batch[1]['coords'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['res_info'] = [pep_batch[1]['res_info'][0][:pep_len]]\n",
    "    pep_batch[1]['seq_lens'] = [[pep_batch[1]['seq_lens'][0][pep_id]]]\n",
    "    pep_batch[1]['coords_loss_mask'] = [pep_batch[1]['coords_loss_mask'][0].narrow(1, 0, pep_len), pep_batch[1]['coords_loss_mask'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['coords_to_seq'] = [pep_batch[1]['coords_to_seq'][0].narrow(1, 0, pep_len), pep_batch[1]['coords_to_seq'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['chain_dict']['begin'] = [pep_batch[1]['chain_dict']['begin'][0].narrow(1, 0, pep_len), pep_batch[1]['chain_dict']['begin'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['chain_dict']['end'] = [pep_batch[1]['chain_dict']['end'][0].narrow(1, 0, pep_len), pep_batch[1]['chain_dict']['end'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['chain_dict']['end'][0][0,-1] = pep_batch[1]['chain_dict']['end'][0][0,-2] \n",
    "    pep_batch[1]['chain_dict']['singles'] = [pep_batch[1]['chain_dict']['singles'][0].narrow(1, 0, pep_len), pep_batch[1]['chain_dict']['singles'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['chain_dict']['ids'] = [pep_batch[1]['chain_dict']['ids'][0].narrow(1, 0, pep_len) - pep_id, pep_batch[1]['chain_dict']['ids'][1].narrow(1, 0, pep_len)]\n",
    "    pep_batch[1]['chain_lens'] = [pep_batch[1]['chain_lens'][0].narrow(1, 0, pep_len) - pep_id, pep_batch[1]['chain_lens'][1].narrow(1, 0, pep_len)]\n",
    "    \n",
    "    prot_batch[0]['string_sequence'] = [prot_batch[0]['string_sequence'][0][pep_len+25:]]\n",
    "    prot_batch_cut_pe = prot_batch[0]['pos_embs'][0][0][pep_len+26] - 1\n",
    "    prot_batch_cut_stc = prot_batch[0]['seq_to_coords'][0].narrow(1, pep_len, prot_len)[0][0]\n",
    "    prot_batch[0]['pos_embs'] = [prot_batch[0]['pos_embs'][0].narrow(1, pep_len+25, prot_len+2) - prot_batch_cut_pe, prot_batch[0]['pos_embs'][1].narrow(1, pep_len+25, prot_len+2)]\n",
    "    prot_batch[0]['pos_embs'][0][0][0] = 0\n",
    "    prot_batch[0]['placeholder_mask'] = [prot_batch[0]['placeholder_mask'][0].narrow(1, pep_len+25, prot_len), prot_batch[0]['placeholder_mask'][1].narrow(1, pep_len+25, prot_len)]\n",
    "    prot_batch[0]['seq_loss_mask'] = [prot_batch[0]['seq_loss_mask'][0].narrow(1, pep_len, prot_len), prot_batch[0]['seq_loss_mask'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[0]['seq_to_coords'] = [prot_batch[0]['seq_to_coords'][0].narrow(1, pep_len, prot_len) - prot_batch_cut_stc, prot_batch[0]['seq_to_coords'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['coords'] = [prot_batch[1]['coords'][0].narrow(1, pep_len, prot_len), prot_batch[1]['coords'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['res_info'] = [prot_batch[1]['res_info'][0][pep_len:]]\n",
    "    prot_batch[1]['seq_lens'] = [[prot_batch[1]['seq_lens'][0][prot_id]]]\n",
    "    prot_batch[1]['coords_loss_mask'] = [prot_batch[1]['coords_loss_mask'][0].narrow(1, pep_len, prot_len), prot_batch[1]['coords_loss_mask'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['coords_to_seq'] = [prot_batch[1]['coords_to_seq'][0].narrow(1, pep_len, prot_len) - prot_batch_cut_stc, prot_batch[1]['coords_to_seq'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['chain_dict']['begin'] = [prot_batch[1]['chain_dict']['begin'][0].narrow(1, pep_len, prot_len) - prot_batch_cut_stc, prot_batch[1]['chain_dict']['begin'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['chain_dict']['end'] = [prot_batch[1]['chain_dict']['end'][0].narrow(1, pep_len, prot_len) - prot_batch_cut_stc, prot_batch[1]['chain_dict']['end'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['chain_dict']['singles'] = [prot_batch[1]['chain_dict']['singles'][0].narrow(1, pep_len, prot_len), prot_batch[1]['chain_dict']['singles'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['chain_dict']['ids'] = [prot_batch[1]['chain_dict']['ids'][0].narrow(1, pep_len, prot_len) - prot_id, prot_batch[1]['chain_dict']['ids'][1].narrow(1, pep_len, prot_len)]\n",
    "    prot_batch[1]['chain_lens'] = [prot_batch[1]['chain_lens'][0].narrow(1, pep_len, prot_len) - prot_id, prot_batch[1]['chain_lens'][1].narrow(1, pep_len, prot_len)]\n",
    "\n",
    "    return pep_batch, prot_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10d2bdf0-6d67-410a-8e63-ca3b9e8bd27a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_batches(pep_batch, prot_batch, pep_batch_r, prot_batch_r):\n",
    "    k0 = pep_batch[0].keys()\n",
    "    k1 = pep_batch[1].keys()\n",
    "    for name, test, real in zip(['pep', 'prot'], [pep_batch, prot_batch], [pep_batch_r, prot_batch_r]):\n",
    "        for i in [0,1]:\n",
    "            if i == 0:\n",
    "                for k in k0:\n",
    "                    if k == 'string_sequence' or k == 'pdb_id':\n",
    "                        if real[i][k] != test[i][k]:\n",
    "                            print(name, i, k)\n",
    "                            raise ValueError\n",
    "                    else:\n",
    "                        if (not ((real[i][k][0] == test[i][k][0]).all())) or (not ((real[i][k][1] == test[i][k][1]).all())):\n",
    "                            print(name, i, k)\n",
    "                            raise ValueError\n",
    "            else:\n",
    "                for k in k1:\n",
    "                    if k == 'chain_dict':\n",
    "                        for kk in test[i][k].keys():\n",
    "                            if (not ((real[i][k][kk][0] == test[i][k][kk][0]).all())) or (not ((real[i][k][kk][1] == test[i][k][kk][1]).all())):\n",
    "                                print(name, i, k)\n",
    "                                raise ValueError\n",
    "                    elif k == 'res_info' or k == 'seq_lens':\n",
    "                        if real[i][k] != test[i][k]:\n",
    "                            print(name, i, k)\n",
    "                            raise ValueError\n",
    "                    else:\n",
    "                        if (not ((real[i][k][0] == test[i][k][0]).all())) or (not ((real[i][k][1] == test[i][k][1]).all())):\n",
    "                            print(name, i, k)\n",
    "                            raise ValueError\n",
    "        print(f'made it: {name}!')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}