make_msas.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88871520",
   "metadata": {},
   "source": [
    "Victor Mikhaylov, vmikhayl@ias.edu<br>Institute for Advanced Study, 2021-2022"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "50dd161d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MHC loading from MHC.pckl. To update the pickle file, set use_pickle to False\n",
      "loaded 26122 MHC sequences in  4.5 s\n",
      "TCR loading from TCR.pckl. To update the pickle file, set use_pickle to False\n",
      "loaded 2835 TCR sequences in  0.2 s\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "cwd=os.getcwd()\n",
    "\n",
    "import tfold.utils.utils as utils\n",
    "import tfold.utils.seq_tools as seq_tools\n",
    "seq_tools.load_mhcs()\n",
    "seq_tools.load_tcrs()\n",
    "\n",
    "from tfold.modeling import result_parse_tools\n",
    "\n",
    "import importlib\n",
    "if 0:\n",
    "    importlib.reload(utils)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "867474f2",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# make MHC MSA index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "12fef02b",
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "896\n"
     ]
    }
   ],
   "source": [
    "with open('./data/experimental_structures/processed_updated/pmhcs.pckl','rb') as f:\n",
    "    pmhcs=pickle.load(f)\n",
    "print(len(pmhcs))\n",
    "#752->896"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c832253c",
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I_A 5\n",
      "745\n",
      "19\n",
      "1\n",
      "2\n",
      "1\n",
      "II_A 7\n",
      "18\n",
      "70\n",
      "19\n",
      "11\n",
      "7\n",
      "2\n",
      "1\n",
      "II_B 4\n",
      "6\n",
      "103\n",
      "14\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "#get MHC unique pdbnums\n",
    "pdbnums={'I_A':{},'II_A':{},'II_B':{}}\n",
    "for x in pmhcs:\n",
    "    if x['class']=='I': \n",
    "        mhc=seq_tools.load_NUMSEQ(x['M'])\n",
    "        pdbnums['I_A'].setdefault('|'.join(mhc.data['pdbnum']),[]).append((mhc.info['locus'],mhc.info['allele'],mhc.seq()))\n",
    "    else:\n",
    "        mhc=seq_tools.load_NUMSEQ(x['M'])\n",
    "        pdbnums['II_A'].setdefault('|'.join(mhc.data['pdbnum']),[]).append((mhc.info['locus'],mhc.info['allele'],mhc.seq()))\n",
    "        mhc=seq_tools.load_NUMSEQ(x['N'])\n",
    "        pdbnums['II_B'].setdefault('|'.join(mhc.data['pdbnum']),[]).append((mhc.info['locus'],mhc.info['allele'],mhc.seq()))\n",
    "        \n",
    "for k,v in pdbnums.items():\n",
    "    print(k,len(v))\n",
    "    for x in v.values():\n",
    "        print(len(x))\n",
    "        \n",
    "#make a lookup table #cl_chain: pdbnum: (id, seq_example)\n",
    "mhc_msa_index={}\n",
    "for cl_chain,v in pdbnums.items():\n",
    "    mhc_msa_index[cl_chain]={}\n",
    "    for i,pdbnum in enumerate(v):\n",
    "        seq=v[pdbnum][0][-1]\n",
    "        mhc_msa_index[cl_chain][pdbnum]=(i,seq)\n",
    "        \n",
    "with open('./data/msas/mhc_msa_index.pckl','wb') as f:\n",
    "    pickle.dump(mhc_msa_index,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aa5cc6d1",
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#seq, name\n",
    "inputs=[]\n",
    "for cl_chain,v in mhc_msa_index.items():\n",
    "    for i,seq in v.values():\n",
    "        inputs.append([seq,'{:s}_{:d}'.format(cl_chain,i)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d8a7f01e",
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['GSHSMRYFYTSVSRPGRGEPRFIAVGYVDDTQFVRFDSDAASQRMEPRAPWIEQEGPEYWDQETRNVKAQSQTDRVDLGTLRGYYNQSEDGSHTIQIMYGCDVGPDGRFLRGYRQDAYDGKDYIALNEDLRSWTAADMAAQITKRKWEAAHAAEQQRAYLEGRCVEWLRRYLENGKETLQRT',\n",
       " 'I_A_0']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ae6f5e5",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# run alignments for MHC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "03d834b1",
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "with open('./data/msas/mhc_msa_index.pckl','rb') as f:\n",
    "    mhc_msa_index=pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "99adf4cd",
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs: 16\n",
      "n tasks: 16\n"
     ]
    }
   ],
   "source": [
    "#prepare jobs\n",
    "\n",
    "output_dir=cwd+'/data/msas/MHC'\n",
    "\n",
    "input_dir='/scratch/vmikhayl/tmp_mhc_msas'           \n",
    "input_dir_server='/scratch/lustre/vmikhayl/tmp_mhc_msas'   #typhon\n",
    "#input_dir_server=input_dir                                  #helios\n",
    "\n",
    "inputs=[] #[seq,name] where name is class_chain_id\n",
    "for cl_chain,v in mhc_msa_index.items():\n",
    "    for i,seq in v.values():\n",
    "        inputs.append([seq,'{:s}_{:d}'.format(cl_chain,i)])\n",
    "print('inputs:',len(inputs))\n",
    "\n",
    "n_tasks=200 #112#56       #our cluster: 64 nodes * 28 cores\n",
    "max_run_time=60*5 #100    #in minutes\n",
    "qos='short'               #medium \n",
    "exclusive=True            #to have enough memory for hhblist\n",
    "local=False\n",
    "slow_nodes=[]\n",
    "\n",
    "sh_path='./make_mhc_msas.sh'\n",
    "python_path='/data/vmikhayl/alphafold-multimer/alphafold/tfold_MSA_tools.py'\n",
    "argstring=f'msa {output_dir} --tmp_dir {input_dir_server}/fasta_tmp'\n",
    "\n",
    "n_tasks=utils.make_task(inputs,n_tasks,sh_path,python_path,input_dir,input_dir_server=input_dir_server,\n",
    "                        qos=qos,max_run_time=max_run_time,slow_nodes=slow_nodes,exclusive=exclusive,\n",
    "                        local=local,argstring=argstring)\n",
    "print('n tasks:',n_tasks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa45208f",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# run alignments for TRV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "98a20182",
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n tasks: 2\n"
     ]
    }
   ],
   "source": [
    "#prepare jobs\n",
    "\n",
    "output_dir=cwd+'/data/msas/TRV'\n",
    "\n",
    "input_dir='/scratch/vmikhayl/tmp_trv_msas'           \n",
    "#input_dir_server='/scratch/lustre/vmikhayl/tmp_mhc_msas'   #typhon\n",
    "input_dir_server=input_dir                                  #helios\n",
    "\n",
    "inputs=[[seq_tools.tcrs['9606','TRAV1-1','01'].seq(),'TRAV'],\n",
    "        [seq_tools.tcrs['9606','TRBV10-1','01'].seq(),'TRBV']]\n",
    "\n",
    "n_tasks=200 #112#56       #our cluster: 64 nodes * 28 cores\n",
    "max_run_time=60*5 #100    #in minutes\n",
    "qos='short'               #medium \n",
    "exclusive=True            #to have enough memory for hhblist\n",
    "local=False              \n",
    "slow_nodes=[]\n",
    "\n",
    "sh_path='./make_trv_msas.sh'\n",
    "python_path='/data/vmikhayl/alphafold-multimer/alphafold/tfold_MSA_tools.py'\n",
    "argstring=f'msa {output_dir} --tmp_dir {input_dir}/fasta_tmp'\n",
    "\n",
    "n_tasks=utils.make_task(inputs,n_tasks,sh_path,python_path,input_dir,input_dir_server=input_dir_server,\n",
    "                        qos=qos,max_run_time=max_run_time,slow_nodes=slow_nodes,exclusive=exclusive,\n",
    "                        local=local,argstring=argstring)\n",
    "print('n tasks:',n_tasks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85d235d0",
   "metadata": {},
   "source": [
    "# make paired pMHC MSAs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fac62f6",
   "metadata": {},
   "source": [
    "## class I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e9df8abc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MHC loading from MHC.pckl. To update the pickle file, set use_pickle to False\n",
      "loaded 26122 MHC sequences in  1.6 s\n"
     ]
    }
   ],
   "source": [
    "import tfold.nn.nn_predict as nn_predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "7500e8b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "uniref90_hits.a3m 9087\n",
      "mgnify_hits.a3m 111\n",
      "bfd_uniclust_hits.a3m 1308\n"
     ]
    }
   ],
   "source": [
    "#for reference, look at MHC MSA length\n",
    "for filename in os.listdir('./data/msas/MHC/I_A_0'):\n",
    "    if filename.endswith('a3m'):\n",
    "        with open('./data/msas/MHC/I_A_0/'+filename) as f:\n",
    "            s=f.read().split('\\n')\n",
    "        print(filename,len(s)//2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "0b66f87e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "peplen distribution:\n",
      "8      11541\n",
      "9     139966\n",
      "10     38351\n",
      "11      9018\n",
      "12       372\n",
      "13       397\n",
      "14       232\n",
      "15        25\n",
      "Name: pep, dtype: int64\n",
      "assays: 199902, mhcs: 161\n",
      "               pep          mhc_a            kd\n",
      "0       AAAAAAYAAM  (10090, D, b)   7333.347094\n",
      "104835    RVYIYYNH  (10090, D, b)  50000.000000\n",
      "105070   SAFEWHLTF  (10090, D, b)  11329.977020\n",
      "105085   SAHVNGELV  (10090, D, b)     24.178188\n",
      "105087   SAIANRLAL  (10090, D, b)      1.238977\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD4CAYAAAD8Zh1EAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAON0lEQVR4nO3df6jd913H8edr6YaVqRFzlZnkmv4R54q4Ou/aQRVn5zRtxeBfttUNy0aotDJBsPUfRfwnYyjbWLcQaihDWRFXNBtxZdDNiq6aVrtuae0IXV1vUuhqzXQuOLK9/ePedmen99zzvck593vO5zwfELjf8/3knHcO7ZMvn/PjpqqQJM2/V/U9gCRpMgy6JDXCoEtSIwy6JDXCoEtSIy7r64F37dpV+/bt6+vhJWkuPfrooy9U1dJG53oL+r59+3jkkUf6enhJmktJ/mPUObdcJKkRBl2SGmHQJakRBl2SGmHQJakRBl2SGjE26EmOJXk+yRdHnE+SDyY5neTxJG+a/JiSpHG6XKHfCxzY5Pz1wP71P4eAj1z6WJKkrRob9Kp6CHhxkyUHgY/WmoeBnUleN6kBJUndTOKToruBZweOV9dve254YZJDrF3Fs7y8PIGHlqT5cO3hBzlz7jwAu3dezj/edd3EH2MSL4pmg9s2/DVIVXW0qlaqamVpacOvIpCkJp05d55nDt/IM4dvfDnskzaJoK8CeweO9wBnJ3C/kqQtmETQjwPvXH+3y1uAr1XVK7ZbJEnTNXYPPcnHgLcCu5KsAn8EvBqgqo4AJ4AbgNPAN4BbpzWsJGm0sUGvqpvHnC/g9olNJEm6KH5SVJIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqREGXZIaYdAlqRGdgp7kQJKnkpxOctcG538gySeSfD7JqSS3Tn5USdJmxgY9yQ7gbuB64Erg5iRXDi27HXiiqt4IvBX40ySvmfCskqRNdLlCvxo4XVVPV9U3gfuAg0NrCvi+JAFeC7wIXJjopJKkTXUJ+m7g2YHj1fXbBn0IeANwFvgC8J6q+vbwHSU5lOSRJI989atfvciRJUkb6RL0bHBbDR3/MvAY8KPAVcCHknz/K/5S1dGqWqmqlaWlpS2OKknaTJegrwJ7B473sHYlPuhW4P5acxr4MvATkxlRktTFZR3WnAT2J7kCOAPcBNwytOYrwNuAf0jyI8DrgacnOagkzZtrDz/ImXPnAdi98/KpP97YoFfVhSR3AA8AO4BjVXUqyW3r548AfwLcm+QLrG3R3FlVL0xxbkmaeWfOneeZwzdu2+N1uUKnqk4AJ4ZuOzLw81nglyY7miTNn+2+Kh/UKeiSpNGGI76dV+WDDLokXaLt3loZxaBL0kXoc2tlFIMuSR3NytbKKAZdkkYYDDjMZsQHGXRJGmFW9sa7MuiSNGAW98a7MuiSFt6s7413ZdAlLbx521oZxV9BJ0mNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmNMOiS1AiDLkmN6BT0JAeSPJXkdJK7Rqx5a5LHkpxK8veTHVOSNM5l4xYk2QHcDbwdWAVOJjleVU8MrNkJfBg4UFVfSfLDU5pXkjRClyv0q4HTVfV0VX0TuA84OLTmFuD+qvoKQFU9P9kxJUnjdAn6buDZgePV9dsG/Tjwg0k+m+TRJO+c1ICSpG7GbrkA2eC22uB+fgZ4G3A58LkkD1fVl77rjpJDwCGA5eXlrU8rSRqpyxX6KrB34HgPcHaDNZ+qqv+tqheAh4A3Dt9RVR2tqpWqWllaWrrYmSVJG+gS9JPA/iRXJHkNcBNwfGjN3wI/l+SyJN8LXAM8OdlRJUmbGbvlUlUXktwBPADsAI5V1akkt62fP1JVTyb5FPA48G3gnqr64jQHlyR9ty576FTVCeDE0G1Hho7fB7xvcqNJkrbCT4pKUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1wqBLUiMMuiQ1otOvoJOk1lx7+EHOnDsPwO6dl/c8zWQYdEkLYzjizxy+seeJJsugS1oYZ86dby7igwy6pKa1uLUyikGX1JzWt1ZGMeiSmtP61sooBl1SExZpa2UUgy5pbi3q1sooBl3S3FrUrZVRDLqkueLWymgGXdLMc2ulG4Muaea5tdKNQZc0k9xa2TqDLmlmuLVyaQy6pF4Z8cnpFPQkB4APADuAe6rq8Ih1bwYeBn69qv56YlNKaooRn46xQU+yA7gbeDuwCpxMcryqnthg3XuBB6YxqKR2+CLndHS5Qr8aOF1VTwMkuQ84CDwxtO53gI8Db57ohJKa4Iuc09cl6LuBZweOV4FrBhck2Q38GnAdmwQ9ySHgEMDy8vJWZ5U0x7wqn74uv1M0G9xWQ8fvB+6sqm9tdkdVdbSqVqpqZWlpqeOIkqQuulyhrwJ7B473AGeH1qwA9yUB2AXckORCVf3NJIaUJI3XJegngf1JrgDOADcBtwwuqKorXvo5yb3AJ425JG2vsUGvqgtJ7mDt3Ss7gGNVdSrJbevnj0x5RklSB53eh15VJ4ATQ7dtGPKq+q1LH0tSC3xny/byk6KSJsoPDfXHoEu6ZEZ8Nhh0SRfFiM8egy6pMyM+2wy6pM78tOdsM+iSNuU7VeaHQZf0Cm6tzCeDLum7Ag5GfF4ZdGlBeRXeHoMuLShf4GyPQZcWiC9wts2gS41za2VxGHSpQUZ8MRl0qUHujy8mgy41wv1xGXRpjrm1okEGXZpjbq1okEGX5oxbKxrFoEtzwK0VdWHQpRllxLVVBl2aUe6Pa6sMujRD3B/XpTDoUs/cWtGkGHSpB0Zc02DQpW1ixDVtBl3aJr7IqWkz6NIU+SKntpNBlybMrRX1xaBLE2DENQsMunSRjLhmjUGXtsCIa5YZdGkMI6550SnoSQ4AHwB2APdU1eGh878B3Ll++HXgt6vq85McVNpORlzzaGzQk+wA7gbeDqwCJ5Mcr6onBpZ9Gfj5qvqvJNcDR4FrpjGwNC1GXPOuyxX61cDpqnoaIMl9wEHg5aBX1T8NrH8Y2DPJIaXt4Ad/NO+6BH038OzA8SqbX32/C/i7jU4kOQQcAlheXu44ojQ9fvBHLekS9GxwW224MPkF1oL+sxudr6qjrG3HsLKysuF9SNM0GHBwa0Vt6RL0VWDvwPEe4OzwoiQ/BdwDXF9V/zmZ8aRL5964FkWXoJ8E9ie5AjgD3ATcMrggyTJwP/COqvrSxKeUtsiIaxGNDXpVXUhyB/AAa29bPFZVp5Lctn7+CPCHwA8BH04CcKGqVqY3tvRKRlyLrtP70KvqBHBi6LYjAz+/G3j3ZEeTtsZ3qWjR+UlRzTXfpSJ9h0HX3HFrRdqYQddcMOLSeAZdM8uIS1tj0DVTjLh08Qy6emfEpckw6OqFEZcmz6Br2xhxaboMuqbKiEvbx6Br4oy41A+Drokw4lL/DLoumhGXZotB15YYcWl2GXRtaPg3+7zEiEuzy6DrZV59S/PNoC84Iy61w6AvICMutcmgLwgjLrXPoDdq+EVNIy61z6DPOd+NIuklBn0OuX0iaSMGfU4YcUnjGPQZZsQlbYVBnwHug0uaBIPeE6++JU2aQZ8yr74lbReDPiGGW1LfDPoWGW5Js8qgDxgV60GGW9KsWsige5UtqUXNBd2rbEmLai6Dvlm0jbWkRTWXQT9z7rzRlqQhr+qyKMmBJE8lOZ3krg3OJ8kH188/nuRNkx9VkrSZsUFPsgO4G7geuBK4OcmVQ8uuB/av/zkEfGTCc0qSxuhyhX41cLqqnq6qbwL3AQeH1hwEPlprHgZ2JnndhGeVJG2iyx76buDZgeNV4JoOa3YDzw0uSnKItSt4gK8neWpL037HrryXFy7y77ZiF/gc4HPgczCnz0Hee9F/9cdGnegS9GxwW13EGqrqKHC0w2NuPlDySFWtXOr9zDOfA58D8DkAn4NBXbZcVoG9A8d7gLMXsUaSNEVdgn4S2J/kiiSvAW4Cjg+tOQ68c/3dLm8BvlZVzw3fkSRpesZuuVTVhSR3AA8AO4BjVXUqyW3r548AJ4AbgNPAN4BbpzcyMIFtmwb4HPgcgM8B+By8LFWv2OqWJM2hTh8skiTNPoMuSY2Yu6CP+xqC1iU5luT5JF/se5a+JNmb5DNJnkxyKsl7+p5pOyX5niT/kuTz6//+P+57pr4k2ZHk35J8su9ZZsFcBb3j1xC07l7gQN9D9OwC8HtV9QbgLcDtC/bfwf8B11XVG4GrgAPr7y5bRO8Bnux7iFkxV0Gn29cQNK2qHgJe7HuOPlXVc1X1r+s//w9r/0Pv7neq7bP+FRtfXz989fqfhXt3Q5I9wI3APX3PMivmLeijvmJACyrJPuCngX/ueZRttb7V8BjwPPDpqlqof/+69wO/D3y75zlmxrwFvdNXDGgxJHkt8HHgd6vqv/ueZztV1beq6irWPpV9dZKf7HmkbZXkV4Dnq+rRvmeZJfMWdL9iQAAkeTVrMf/Lqrq/73n6UlXngM+yeK+rXAv8apJnWNt6vS7JX/Q7Uv/mLehdvoZAjUsS4M+BJ6vqz/qeZ7slWUqyc/3ny4FfBP6916G2WVX9QVXtqap9rHXgwar6zZ7H6t1cBb2qLgAvfQ3Bk8BfVdWpfqfaXkk+BnwOeH2S1STv6numHlwLvIO1q7LH1v/c0PdQ2+h1wGeSPM7aRc6nq8q37cmP/ktSK+bqCl2SNJpBl6RGGHRJaoRBl6RGGHRJaoRBl6RGGHRJasT/A2cNPv1a20/0AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "after reducing by true kd at threshold 100.0:\n",
      "assays: 33365, mhcs: 141\n",
      "making Kd predictions for 33365 pmhcs...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vmikhayl/anaconda3/envs/bio-p310/lib/python3.10/site-packages/keras/engine/functional.py:559: UserWarning: Input dict contained keys ['regmask', 'tails'] which did not match any model input. They will be ignored by the model.\n",
      "  inputs = self._flatten_to_reference_inputs(inputs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "seqnn run completed in 161.24462914466858\n",
      "after reducing by predicted kd at threshold 100.0:\n",
      "assays: 27763, mhcs: 137\n",
      "after subsampling 100 assays per mhc:\n",
      "assays: 8232, mhcs: 137\n"
     ]
    }
   ],
   "source": [
    "#read pmhc assays\n",
    "train=pd.read_pickle('./data/pmhc_assays/train_i.pckl') \n",
    "print('peplen distribution:')\n",
    "print(train['pep'].map(len).value_counts().sort_index()) #len distribution\n",
    "train=train.sort_values('mhc_a')\n",
    "print('assays: {}, mhcs: {}'.format(len(train),len(train['mhc_a'].unique())))\n",
    "print(train.head(5))\n",
    "plt.figure()\n",
    "plt.hist(np.log10(train['kd'].values),bins=100,histtype='step',cumulative=1,density=True)\n",
    "plt.show()\n",
    "#note: in case AF relies on seq order when subsampling the MSA, sort entries by MHC; \n",
    "#(I believe, it uses the order only during training, not in inference)\n",
    "\n",
    "#reduce to good binders\n",
    "kd_threshold=100.\n",
    "train=train[train['kd']<kd_threshold]\n",
    "print('after reducing by true kd at threshold {}:'.format(kd_threshold))\n",
    "print('assays: {}, mhcs: {}'.format(len(train),len(train['mhc_a'].unique())))\n",
    "#predict kds\n",
    "t0=time.time()\n",
    "train=nn_predict.predict(train,'I',mhc_as_obj=False)\n",
    "print('seqnn run completed in {}'.format(time.time()-t0))\n",
    "#reduce to predicted good binders\n",
    "train=train[train['seqnn_logkd']<np.log10(kd_threshold)]\n",
    "print('after reducing by predicted kd at threshold {}:'.format(kd_threshold))\n",
    "print('assays: {}, mhcs: {}'.format(len(train),len(train['mhc_a'].unique())))\n",
    "\n",
    "#choose a subset\n",
    "n_per_mhc=100\n",
    "mhc_u=train['mhc_a'].unique()\n",
    "train1=[]\n",
    "for m in mhc_u:\n",
    "    df=train[train['mhc_a']==m]\n",
    "    train1.append(df.sample(min(n_per_mhc,len(df))))\n",
    "train=pd.concat(train1)\n",
    "print('after subsampling {} assays per mhc:'.format(n_per_mhc))\n",
    "print('assays: {}, mhcs: {}'.format(len(train),len(train['mhc_a'].unique())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "0670a0da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total pep pdbnums: 51\n",
      "['   09', '   1 ', '   2 ', '   3 ', '   4 ', '   5 ', '   51', '   52', '   53', '   54', '   57', '   58', '   59', '   6 ', '   7 ', '   8 ', '   9 ', '  10 ', '  11 ', '  12 ', '  13 ', '  14 ', '  15 ', '  16 ']\n",
      "different cl I pdbnums: 13\n",
      "res positions in cl I pdbnums: 186\n",
      "['   1 ' '   2 ' '   3 ' '   4 ' '   5 ' '   6 ' '   7 ' '   8 ' '   9 '\n",
      " '  10 ' '  11 ' '  12 ' '  13 ' '  14 ' '  15 ' '  16 ' '  17 ' '  18 '\n",
      " '  19 ' '  20 ' '  21 ' '  22 ' '  23 ' '  24 ' '  25 ' '  26 ' '  27 '\n",
      " '  28 ' '  29 ' '  30 ' '  31 ' '  32 ' '  33 ' '  34 ' '  35 ' '  36 '\n",
      " '  37 ' '  38 ' '  39 ' '  40 ' '  41 ' '  42 ' '  43 ' '  44 ' '  45 '\n",
      " '  46 ' '  47 ' '  48 ' '  49 ' '  50 ' '  51 ' '  52 ' '  53 ' '  54 '\n",
      " '  54A' '  54B' '  54C' '  55 ' '  56 ' '  57 ' '  58 ' '  59 ' '  60 '\n",
      " '  61 ' '  62 ' '  63 ' '  64 ' '  65 ' '  66 ' '  67 ' '  68 ' '  69 '\n",
      " '  70 ' '  71 ' '  72 ' '  73 ' '  74 ' '  75 ' '  76 ' '  77 ' '  78 '\n",
      " '  79 ' '  80 ' '  81 ' '  82 ' '  83 ' '  84 ' '  85 ' '  86 ' '  87 '\n",
      " '  88 ' '  89 ' '  90 ' '1001 ' '1002 ' '1003 ' '1004 ' '1005 ' '1006 '\n",
      " '1007 ' '1008 ' '1009 ' '1010 ' '1011 ' '1012 ' '1013 ' '1014 ' '1015 '\n",
      " '1016 ' '1017 ' '1018 ' '1019 ' '1020 ' '1021 ' '1022 ' '1023 ' '1024 '\n",
      " '1025 ' '1026 ' '1027 ' '1028 ' '1029 ' '1030 ' '1031 ' '1032 ' '1033 '\n",
      " '1034 ' '1035 ' '1036 ' '1037 ' '1038 ' '1039 ' '1042 ' '1043 ' '1044 '\n",
      " '1045 ' '1046 ' '1047 ' '1048 ' '1049 ' '1050 ' '1051 ' '1052 ' '1053 '\n",
      " '1054 ' '1055 ' '1056 ' '1057 ' '1058 ' '1059 ' '1060 ' '1061 ' '1061A'\n",
      " '1061B' '1062 ' '1063 ' '1064 ' '1065 ' '1066 ' '1067 ' '1068 ' '1069 '\n",
      " '1070 ' '1071 ' '1072 ' '1072A' '1073 ' '1074 ' '1075 ' '1076 ' '1077 '\n",
      " '1078 ' '1079 ' '1080 ' '1081 ' '1082 ' '1083 ' '1084 ' '1085 ' '1086 '\n",
      " '1087 ' '1088 ' '1089 ' '1090 ' '1091 ' '1092 ']\n"
     ]
    }
   ],
   "source": [
    "#make master pdbnums for mhc and pep\n",
    "\n",
    "###pep###\n",
    "import tfold.nn.nn_utils as nn_utils\n",
    "import tfold.modeling.template_tools as template_tools\n",
    "pep_pdbnum=set()\n",
    "pep_pdbnums_all=set()\n",
    "for l in range(8,16):\n",
    "    for tails in nn_utils.generate_registers_I(l):\n",
    "        c_pdbnum=template_tools._make_pep_pdbnums_I(l,tails[0],tails[1])\n",
    "        pep_pdbnums_all.add(tuple(c_pdbnum))\n",
    "        pep_pdbnum|=set(c_pdbnum)\n",
    "pep_pdbnums_all=[list(x) for x in pep_pdbnums_all]\n",
    "pep_pdbnum=sorted(list(pep_pdbnum))\n",
    "print('total pep pdbnums:',len(pep_pdbnums_all))\n",
    "print(pep_pdbnum)\n",
    "\n",
    "###mhc###\n",
    "#NOTE: in MHC MSAs, used only pdbnums for mhcs in our structure database; here using all (for cl I)\n",
    "#(in the future: make MHC MSAs for all pdbnums? will be a lot of them for cl II)\n",
    "\n",
    "#collect all cl I mhc pdbnums\n",
    "mhc_pdbnums_all=set()\n",
    "for mhc in seq_tools.mhcs.values():\n",
    "    if mhc.info['class']=='I':\n",
    "        mhc_pdbnums_all.add(tuple(mhc.data['pdbnum']))\n",
    "mhc_pdbnums_all=[list(x) for x in mhc_pdbnums_all]\n",
    "print('different cl I pdbnums:',len(mhc_pdbnums_all))\n",
    "\n",
    "#all res positions in them\n",
    "pdbnum_entries=set()\n",
    "for pdbnum in mhc_pdbnums_all:\n",
    "    pdbnum_entries|=set(pdbnum)\n",
    "pdbnum_entries=np.sort(list(pdbnum_entries))\n",
    "print('res positions in cl I pdbnums:',len(pdbnum_entries))\n",
    "\n",
    "#permute to ensure order\n",
    "for pdbnum in mhc_pdbnums_all:\n",
    "    pdbnum_entries[np.isin(pdbnum_entries,pdbnum)]=pdbnum\n",
    "#check\n",
    "for pdbnum in mhc_pdbnums_all:\n",
    "    ind=np.isin(pdbnum_entries,pdbnum)\n",
    "    if np.sum(ind)!=len(pdbnum):\n",
    "        print('oops')\n",
    "    if '|'.join(pdbnum_entries[ind])!='|'.join(pdbnum):\n",
    "        print('oops')\n",
    "mhc_pdbnum=np.array(pdbnum_entries)\n",
    "print(mhc_pdbnum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "19f8f029",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(8232, 24) (8232, 186)\n"
     ]
    }
   ],
   "source": [
    "#convert peps and mhcs in assays to NUMSEQ, extend to master pdbnum\n",
    "def pep_to_numseq(x):\n",
    "    seq=x['pep']\n",
    "    tails=x['seqnn_tails']\n",
    "    return seq_tools.NUMSEQ(seq=seq,pdbnum=template_tools._make_pep_pdbnums_I(len(seq),tails[0],tails[1]))\n",
    "def mhc_to_numseq(x):\n",
    "    return seq_tools.mhcs[x['mhc_a']]\n",
    "def extend_numseq(s,pdbnum):\n",
    "    s_ext=seq_tools.NUMSEQ(seq='_'*len(pdbnum),pdbnum=pdbnum) #use '_' for gap symbol, to replace later\n",
    "    ind1=np.isin(pdbnum,s.data['pdbnum'])\n",
    "    ind2=np.isin(s.data['pdbnum'],pdbnum)\n",
    "    s_ext.data[ind1]=s.data[ind2]\n",
    "    return s_ext.data['seq']\n",
    "pep_assays=train.apply(pep_to_numseq,axis=1)\n",
    "pep_assays=np.stack(pep_assays.map(lambda x: extend_numseq(x,pep_pdbnum)).values)\n",
    "mhc_assays=train.apply(mhc_to_numseq,axis=1)\n",
    "mhc_assays=np.stack(mhc_assays.map(lambda x: extend_numseq(x,mhc_pdbnum)).values)\n",
    "print(pep_assays.shape,mhc_assays.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "3be4d512",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "completed in 20.645625114440918\n"
     ]
    }
   ],
   "source": [
    "#reduce to a3m, save pdbnum->filename index\n",
    "\n",
    "pmhc_msa_index={'I_P':{},'I_M':{}} #cl_chain: pdbnum: id        \n",
    "\n",
    "lower_vect=np.vectorize(lambda s:s.lower())\n",
    "repl_gap_vect=np.vectorize(lambda s:s.replace('_','-'))\n",
    "def reduce_matrix_to_pdbnum(matrix,ind):\n",
    "    '''\n",
    "    entries within pdbnum are kept, with gap symbols replaced by '-';\n",
    "    entries outside pdbnum are made lowercase, with gap symbols dropped;\n",
    "    outputs an array of strings\n",
    "    '''\n",
    "    matrix=matrix.copy()\n",
    "    matrix[:,~ind]=lower_vect(matrix[:,~ind])\n",
    "    matrix[:,ind]=repl_gap_vect(matrix[:,ind])\n",
    "    return matrix\n",
    "\n",
    "t0=time.time()\n",
    "#pep\n",
    "for i,pdbnum in enumerate(pep_pdbnums_all):\n",
    "    ind=np.isin(pep_pdbnum,pdbnum) #which to keep\n",
    "    pep_assays_reduced=reduce_matrix_to_pdbnum(pep_assays,ind)\n",
    "    pep_a3m=[]\n",
    "    for j,line in enumerate(pep_assays_reduced):\n",
    "        pep_a3m.append(f'>{j}')\n",
    "        pep_a3m.append(''.join(line).replace('_',''))        \n",
    "    with open(f'./data/msas/pMHC/I_P_{i}.a3m','w') as f:\n",
    "        f.write('\\n'.join(pep_a3m))\n",
    "    pmhc_msa_index['I_P']['|'.join(pdbnum)]=i\n",
    "    \n",
    "#mhc\n",
    "for i,pdbnum in enumerate(mhc_pdbnums_all):\n",
    "    ind=np.isin(mhc_pdbnum,pdbnum) #which to keep\n",
    "    mhc_assays_reduced=reduce_matrix_to_pdbnum(mhc_assays,ind)\n",
    "    mhc_a3m=[]\n",
    "    for j,line in enumerate(mhc_assays_reduced):\n",
    "        mhc_a3m.append(f'>{j}')\n",
    "        mhc_a3m.append(''.join(line).replace('_',''))        \n",
    "    with open(f'./data/msas/pMHC/I_M_{i}.a3m','w') as f:\n",
    "        f.write('\\n'.join(mhc_a3m))\n",
    "    pmhc_msa_index['I_M']['|'.join(pdbnum)]=i\n",
    "\n",
    "#save lookup table\n",
    "with open('./data/msas/pmhc_msa_index.pckl','wb') as f:\n",
    "    pickle.dump(pmhc_msa_index,f)\n",
    "print('completed in {}'.format(time.time()-t0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53713e06",
   "metadata": {},
   "source": [
    "## class II"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d4582086",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading...\n",
      "completed in 19.41299533843994\n",
      "688216\n",
      "Index(['pmhc_id', 'model_id', 'tpl_tails', 'best_score', 'best_mhc_score',\n",
      "       'register_identified', 'af_tails', 'pep_lddt', 'mhc_lddt', 'class',\n",
      "       'pep', 'mhc_a', 'mhc_b', 'kd', 'n_reg_filtered', 'af_n_reg',\n",
      "       'seqnn_logkd'],\n",
      "      dtype='object')\n",
      "reducing to best...\n",
      "688216\n",
      "42413\n",
      "completed in 90.40511631965637\n"
     ]
    }
   ],
   "source": [
    "#read modeling results\n",
    "result_dir=f'/scratch/vmikhayl/pmhc_modeling/pmhc_assays_train_II_dataset/processed'\n",
    "print('loading...')\n",
    "t0=time.time()\n",
    "results_df=[]\n",
    "for x in os.listdir(result_dir):\n",
    "    results_df.append(pd.read_pickle(result_dir+f'/{x}/result.pckl'))\n",
    "results_df=pd.concat(results_df)\n",
    "print('completed in {}'.format(time.time()-t0))\n",
    "print(len(results_df))\n",
    "print(results_df.columns)\n",
    "\n",
    "#reduce to best\n",
    "print('reducing to best...')\n",
    "t0=time.time()\n",
    "results_df['lddt_score']=results_df['pep_lddt'].map(result_parse_tools.lddt_score)\n",
    "reduced_df=result_parse_tools.reduce_to_best(results_df,['pmhc_id'],'lddt_score',how='min')\n",
    "print(len(results_df))\n",
    "print(len(reduced_df))\n",
    "print('completed in {}'.format(time.time()-t0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "db978ead",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYQAAAD4CAYAAADsKpHdAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAYg0lEQVR4nO3df6zddZ3n8edLYLRZB1AobG3LlF06yQCJVW663ZBsWHHXRs0UExhrMkJ2mq1LMGLGzUj9Y2XWNIHsCrPowKYOhMKI0KAuDUvXZUBj3NTixUWhIGt36ci1Da2AWLPC2s57/zifDofL6b3n/jynvc9HcnK+5/39fr79fL9w7/t+fnzPJ1WFJElvGXQFJEnDwYQgSQJMCJKkxoQgSQJMCJKk5uRBV2C6zjzzzFqxYsWgqyFJx5XHH3/8F1W1uNe+4zYhrFixgtHR0UFXQ5KOK0n+9lj77DKSJAEmBElSY0KQJAEmBElSY0KQJAEmBElSY0KQJAEmBElSY0KQJAHH8ZPKkubfxTc8ys9/+ZtBV2PBW3r6Iv7Hde+b9fOaEKQFYLZ+kS89fRF7b/jQLNRIM7Hiuv86J+c1IUjHoan+gvcXufphQpCGTD+/7P0Fr7lgQpAG5Fi/+P1lr0ExIUjzqDsJ+Itfw8aEIM2yibp8TAIaZiYEaRb4l79OBH0nhCQnAaPAz6vqw0neCdwHrAD2An9UVS+3YzcBG4AjwKeq6lstfhFwJ7AIeAi4tqoqyVuBu4CLgBeBj1bV3lm4PmnOmAR0oplKC+Fa4Bng1Pb5OuCRqrohyXXt82eTnA+sBy4A3gX8TZLfr6ojwG3ARuD7dBLCWmAHneTxclWdl2Q9cCPw0RlfnTTLTAI6kfWVEJIsAz4EbAb+tIXXAZe07a3Ad4DPtvi9VfUa8FySPcDqJHuBU6tqZzvnXcBldBLCOuD6dq77gS8nSVXV9C9Nmh0mAS0U/bYQ/gL4M+B3u2JnV9V+gKran+SsFl9KpwVw1FiL/bZtj48fLfN8O9fhJK8AZwC/6K5Eko10Whicc845fVZdmpmf//I3JgEtCJMmhCQfBg5U1eNJLunjnOkRqwniE5V5Y6BqC7AFYGRkxNaD5sz4VoG0EPTTQrgY+MMkHwTeBpya5K+BF5Isaa2DJcCBdvwYsLyr/DJgX4sv6xHvLjOW5GTgNOClaV6TNC12DWmhmzQhVNUmYBNAayH826r64yT/AbgKuKG9P9CKbAfuSXITnUHllcBjVXUkyaEka4BdwJXAl7rKXAXsBC4HHnX8QPPNriEtdDN5DuEGYFuSDcDPgCsAqmp3km3A08Bh4Jo2wwjgal6fdrqjvQBuB+5uA9Av0ZmlJM05u4ak100pIVTVd+jMJqKqXgQuPcZxm+nMSBofHwUu7BF/lZZQpPlkq0B6nU8qa8GxVSD1ZkLQgmOrQOrNhKAFwVaBNDkTghYEWwXS5EwIOmHZKpCmxoSgE5atAmlq3jLoCkiShoMtBJ1Q7CaSps+EoBOK3UTS9NllJEkCbCHoBGA3kTQ7TAg67tlNJM0Ou4wkSYAtBB2n7CaSZp8JQcclu4mk2WeXkSQJ6KOFkORtwHeBt7bj76+qzye5HvjXwMF26Oeq6qFWZhOwATgCfKqqvtXiF/H6imkPAddWVSV5K3AXcBHwIvDRqto7S9eoE4TdRNLc6qfL6DXgfVX16ySnAN9LcnTpy5ur6j92H5zkfDpLYF5AZ03lv0ny+20ZzduAjcD36SSEtXSW0dwAvFxV5yVZD9wIfHTml6cTid1E0tyatMuoOn7dPp7SXjVBkXXAvVX1WlU9B+wBVidZApxaVTurqui0CC7rKrO1bd8PXJokU74aSdK09TWGkOSkJE8AB4CHq2pX2/XJJD9OckeSd7TYUuD5ruJjLba0bY+Pv6FMVR0GXgHO6FGPjUlGk4wePHhw/G5J0gz0NcuodfesSnI68M0kF9Lp/vkCndbCF4AvAn8C9PrLviaIM8m+7npsAbYAjIyMTNRK0QnCcQNp/kxp2mlV/TLJd4C13WMHSb4CPNg+jgHLu4otA/a1+LIe8e4yY0lOBk4DXppK3XRictxAmj+TdhklWdxaBiRZBLwf+EkbEzjqI8BTbXs7sD7JW5OcC6wEHquq/cChJGva+MCVwANdZa5q25cDj7ZxBknSPOmnhbAE2JrkJDoJZFtVPZjk7iSr6HTt7AU+AVBVu5NsA54GDgPXtC4ngKt5fdrpjvYCuB24O8keOi2D9TO/NB2v7CaSBmPShFBVPwbe0yP+8QnKbAY294iPAhf2iL8KXDFZXbQw2E0kDYZPKkuSABOCJKkxIUiSAL/tVEPCgWRp8EwIGgoOJEuDZ5eRJAkwIUiSGruMNDCOG0jDxYSggXHcQBoudhlJkgATgiSpMSFIkgDHEDTPHEiWhpcJQfPKgWRpeNllJEkC+lsx7W1JHkvyoyS7k/x5i78zycNJftre39FVZlOSPUmeTfKBrvhFSZ5s+25pK6fRVle7r8V3JVkxB9cqSZpAPy2E14D3VdW7gVXA2iRrgOuAR6pqJfBI+0yS8+mseHYBsBa4ta22BnAbsJHOspor236ADcDLVXUecDNw48wvTZI0Ff2smFbAr9vHU9qrgHXAJS2+FfgO8NkWv7eqXgOea8tirk6yFzi1qnYCJLkLuIzOMprrgOvbue4Hvpwkrqt8YnAgWTo+9DWo3P7Cfxw4D/jLqtqV5Oyq2g9QVfuTnNUOXwp8v6v4WIv9tm2Pjx8t83w71+EkrwBnAL8YV4+NdFoYnHPOOf1eowbMgWTp+NDXoHJVHamqVcAyOn/tv2ld5C7pdYoJ4hOVGV+PLVU1UlUjixcvnqTWkqSpmNIso6r6JZ2uobXAC0mWALT3A+2wMWB5V7FlwL4WX9Yj/oYySU4GTgNemkrdJEkz088so8VJTm/bi4D3Az8BtgNXtcOuAh5o29uB9W3m0Ll0Bo8fa91Lh5KsabOLrhxX5ui5LgcedfxAkuZXP2MIS4CtbRzhLcC2qnowyU5gW5INwM+AKwCqaneSbcDTwGHgmqo60s51NXAnsIjOYPKOFr8duLsNQL9EZ5aSJGke9TPL6MfAe3rEXwQuPUaZzcDmHvFR4E3jD1X1Ki2h6MTgzCLp+ONXV2hOOLNIOv741RWSJMCEIElqTAiSJMCEIElqHFTWrHFmkXR8MyFo1jizSDq+2WUkSQJMCJKkxoQgSQJMCJKkxoQgSQKcZaQZ6J5mCk41lY53JgRNm9NMpROLXUaSJMCEIElq+llCc3mSbyd5JsnuJNe2+PVJfp7kifb6YFeZTUn2JHk2yQe64hclebLtu6UtpUlbbvO+Ft+VZMUcXKskaQL9tBAOA5+pqj8A1gDXJDm/7bu5qla110MAbd964AJgLXBrW34T4DZgI511lle2/QAbgJer6jzgZuDGmV+aJGkqJk0IVbW/qn7Ytg8BzwBLJyiyDri3ql6rqueAPcDqJEuAU6tqZ1UVcBdwWVeZrW37fuDSo60HSdL8mNIYQuvKeQ+wq4U+meTHSe5I8o4WWwo831VsrMWWtu3x8TeUqarDwCvAGT3+/Y1JRpOMHjx4cCpVlyRNou9pp0neDnwd+HRV/SrJbcAXgGrvXwT+BOj1l31NEGeSfa8HqrYAWwBGRkbetF9zz6+4lk5cfSWEJKfQSQZfrapvAFTVC137vwI82D6OAcu7ii8D9rX4sh7x7jJjSU4GTgNemurFaO757IF04upnllGA24FnquqmrviSrsM+AjzVtrcD69vMoXPpDB4/VlX7gUNJ1rRzXgk80FXmqrZ9OfBoG2eQJM2TfloIFwMfB55M8kSLfQ74WJJVdLp29gKfAKiq3Um2AU/TmaF0TVUdaeWuBu4EFgE72gs6CefuJHvotAzWz+SiJElTN2lCqKrv0buP/6EJymwGNveIjwIX9oi/ClwxWV0kSXPHJ5UlSYAJQZLUmBAkSYBff60++OyBtDCYEDQpnz2QFga7jCRJgAlBktSYECRJgAlBktSYECRJgAlBktQ47VQ9+eyBtPCYENSTzx5IC49dRpIkwIQgSWr6WTFteZJvJ3kmye4k17b4O5M8nOSn7f0dXWU2JdmT5NkkH+iKX5TkybbvlrZyGm11tftafFeSFXNwrZKkCfTTQjgMfKaq/gBYA1yT5HzgOuCRqloJPNI+0/atBy4A1gK3Jjmpnes2YCOdZTVXtv0AG4CXq+o84Gbgxlm4NknSFEyaEKpqf1X9sG0fAp4BlgLrgK3tsK3AZW17HXBvVb1WVc8Be4DVbQ3mU6tqZ1sv+a5xZY6e637g0qOtB0nS/JjSGELrynkPsAs4u6r2QydpAGe1w5YCz3cVG2uxpW17fPwNZarqMPAKcEaPf39jktEkowcPHpxK1SVJk+g7ISR5O/B14NNV9auJDu0RqwniE5V5Y6BqS1WNVNXI4sWLJ6uyJGkK+noOIckpdJLBV6vqGy38QpIlVbW/dQcdaPExYHlX8WXAvhZf1iPeXWYsycnAacBL07gezYAPo0kL26QJofXl3w48U1U3de3aDlwF3NDeH+iK35PkJuBddAaPH6uqI0kOJVlDp8vpSuBL4861E7gceLSNM2ge+TCatLD100K4GPg48GSSJ1rsc3QSwbYkG4CfAVcAVNXuJNuAp+nMULqmqo60clcDdwKLgB3tBZ2Ec3eSPXRaButndlmSpKmaNCFU1ffo3ccPcOkxymwGNveIjwIX9oi/SksokqTB8EllSRJgQpAkNSYESRJgQpAkNSYESRJgQpAkNa6YtsD5dLKko0wIC5xPJ0s6yi4jSRJgQpAkNSYESRJgQpAkNSYESRJgQpAkNSYESRLQR0JIckeSA0me6opdn+TnSZ5orw927duUZE+SZ5N8oCt+UZIn275b2kpsJHlrkvtafFeSFbN8jZKkPvTTQrgTWNsjfnNVrWqvhwCSnE9ntbMLWplbk5zUjr8N2EhnSc2VXefcALxcVecBNwM3TvNaJEkz0M+Kad+dwl/t64B7q+o14Lm2JObqJHuBU6tqJ0CSu4DL6CyhuQ64vpW/H/hykrim8tzx6yok9TKTr674ZJIrgVHgM1X1MrAU+H7XMWMt9tu2PT5Oe38eoKoOJ3kFOAP4xQzqpgn4dRWSepnuoPJtwD8GVgH7gS+2eK+1l2uC+ERl3iTJxiSjSUYPHjw4pQpLkiY2rYRQVS9U1ZGq+jvgK8DqtmsMWN516DJgX4sv6xF/Q5kkJwOnAS8d49/dUlUjVTWyePHi6VRdknQM00oISZZ0ffwIcHQG0nZgfZs5dC6dwePHqmo/cCjJmja76Ergga4yV7Xty4FHHT+QpPk36RhCkq8BlwBnJhkDPg9ckmQVna6dvcAnAKpqd5JtwNPAYeCaqjrSTnU1nRlLi+gMJu9o8duBu9sA9Et0ZilJkuZZP7OMPtYjfPsEx28GNveIjwIX9oi/ClwxWT0kSXPLJ5UlSYAJQZLUmBAkSYBrKi8YPp0saTImhAXCp5MlTcYuI0kSYEKQJDUmBEkSYEKQJDUmBEkSYEKQJDUmBEkSYEKQJDU+mHYC8+lkSVNhQjiB+XSypKmwy0iSBPSREJLckeRAkqe6Yu9M8nCSn7b3d3Tt25RkT5Jnk3ygK35RkifbvlvaUpq05Tbva/FdSVbM8jVKkvrQTwvhTmDtuNh1wCNVtRJ4pH0myfl0lsC8oJW5NclJrcxtwEY66yyv7DrnBuDlqjoPuBm4cboXI0mavkkTQlV9l85ax93WAVvb9lbgsq74vVX1WlU9B+wBVidZApxaVTurqoC7xpU5eq77gUuPth4kSfNnumMIZ1fVfoD2flaLLwWe7zpurMWWtu3x8TeUqarDwCvAGb3+0SQbk4wmGT148OA0qy5J6mW2B5V7/WVfE8QnKvPmYNWWqhqpqpHFixdPs4qSpF6mmxBeaN1AtPcDLT4GLO86bhmwr8WX9Yi/oUySk4HTeHMXlSRpjk33OYTtwFXADe39ga74PUluAt5FZ/D4sao6kuRQkjXALuBK4EvjzrUTuBx4tI0zaBp8GE3SdE2aEJJ8DbgEODPJGPB5OolgW5INwM+AKwCqaneSbcDTwGHgmqo60k51NZ0ZS4uAHe0FcDtwd5I9dFoG62flyhYoH0aTNF2TJoSq+tgxdl16jOM3A5t7xEeBC3vEX6UlFEnS4PiksiQJMCFIkhoTgiQJMCFIkhoTgiQJMCFIkhoXyDkB+DCapNlgQjgB+DCapNlgl5EkCTAhSJIaE4IkCTAhSJIaE4IkCTAhSJIap50ep3z2QNJsm1FCSLIXOAQcAQ5X1UiSdwL3ASuAvcAfVdXL7fhNwIZ2/Keq6lstfhGvL57zEHCtq6ZNzGcPJM222egy+udVtaqqRtrn64BHqmol8Ej7TJLz6ayGdgGwFrg1yUmtzG3ARjpLbq5s+yVJ82guxhDWAVvb9lbgsq74vVX1WlU9B+wBVidZApxaVTtbq+CurjKSpHky04RQwH9P8niSjS12dlXtB2jvZ7X4UuD5rrJjLba0bY+Pv0mSjUlGk4wePHhwhlWXJHWb6aDyxVW1L8lZwMNJfjLBsekRqwnibw5WbQG2AIyMjDjGIEmzaEYJoar2tfcDSb4JrAZeSLKkqva37qAD7fAxYHlX8WXAvhZf1iOucZxZJGkuTTshJPkHwFuq6lDb/pfAvwe2A1cBN7T3B1qR7cA9SW4C3kVn8PixqjqS5FCSNcAu4ErgS9Ot14nMmUWS5tJMWghnA99McvQ891TVf0vyA2Bbkg3Az4ArAKpqd5JtwNPAYeCaqjrSznU1r0873dFekqR5NO2EUFX/B3h3j/iLwKXHKLMZ2NwjPgpcON26SJJmzq+ukCQBfnXF0HMgWdJ8MSEMOQeSJc0Xu4wkSYAJQZLUmBAkSYBjCEPJgWRJg2BCGEIOJEsaBLuMJEmACUGS1NhlNCQcN5A0aCaEIeG4gaRBs8tIkgTYQhgou4kkDRMTwgDZTSRpmJgQ5pmtAknDamgSQpK1wH8CTgL+qqpuGHCV5oStAknDaigSQpKTgL8E/gUwBvwgyfaqenqwNZsdtgokHQ+GIiEAq4E9bVlOktwLrKOz/vJQ6/5lfyxLT19kq0DS0BuWhLAUeL7r8xjwT8YflGQjsLF9/HWSZ2e5HmcCv5jlc/K3QDbN2unmpI6zaNjrB9ZxNgx7/WD46zij+uXGaf+7v3esHcOSENIjVm8KVG0BtsxZJZLRqhqZq/PPhmGv47DXD6zjbBj2+sHw13EY6zcsD6aNAcu7Pi8D9g2oLpK0IA1LQvgBsDLJuUl+B1gPbB9wnSRpQRmKLqOqOpzkk8C36Ew7vaOqdg+gKnPWHTWLhr2Ow14/sI6zYdjrB8Nfx6GrX6re1FUvSVqAhqXLSJI0YCYESRKwQBNCkrVJnk2yJ8l1PfZfkuSVJE+017+b5/rdkeRAkqeOsT9Jbmn1/3GS985n/fqs46Dv4fIk307yTJLdSa7tcczA7mOf9Rv0PXxbkseS/KjV8c97HDPIe9hP/QZ6D7vqcVKS/5nkwR77Bv7z/PeqakG96Axa/2/gHwG/A/wIOH/cMZcADw6wjv8MeC/w1DH2fxDYQef5jTXAriGs46Dv4RLgvW37d4H/1eO/88DuY5/1G/Q9DPD2tn0KsAtYM0T3sJ/6DfQedtXjT4F7etVlGH6ej74WYgvh778mo6r+H3D0azKGRlV9F3hpgkPWAXdVx/eB05MsmZ/adfRRx4Gqqv1V9cO2fQh4hs4T8d0Gdh/7rN9Atfvy6/bxlPYaPwtlkPewn/oNXJJlwIeAvzrGIQP/eT5qISaEXl+T0esH8Z+2puiOJBfMT9X61u81DNpQ3MMkK4D30PkLsttQ3McJ6gcDvoetq+MJ4ADwcFUN1T3so34w+P8P/wL4M+DvjrF/KP4/hIWZEPr5mowfAr9XVe8GvgT8l7mu1BT19VUfAzYU9zDJ24GvA5+uql+N392jyLzex0nqN/B7WFVHqmoVnW8PWJ3kwnGHDPQe9lG/gd7DJB8GDlTV4xMd1iM2kJ/nhZgQJv2ajKr61dGmaFU9BJyS5Mz5q+Kkhv6rPobhHiY5hc4v269W1Td6HDLQ+zhZ/YbhHnbV5ZfAd4C143YNxf+Lx6rfENzDi4E/TLKXTvf0+5L89bhjhuIewsJMCJN+TUaSf5gkbXs1nfv04rzX9Ni2A1e22QlrgFeqav+gK9Vt0Pew/du3A89U1U3HOGxg97Gf+g3BPVyc5PS2vQh4P/CTcYcN8h5OWr9B38Oq2lRVy6pqBZ3fNY9W1R+PO2xofp6H4qsr5lMd42sykvybtv8/A5cDVyc5DPwGWF9tOsB8SPI1OrMjzkwyBnyezoDZ0fo9RGdmwh7g/wL/ar7qNoU6DvQe0vnL7OPAk62PGeBzwDlddRzkfeynfoO+h0uAreksYPUWYFtVPTjuZ2WQ97Cf+g36HvY0RPfwjfUagnsjSRoCC7HLSJLUgwlBkgSYECRJjQlBkgSYECRJjQlBkgSYECRJzf8HrNBnZ8mtBoYAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "non-canonical tails dropped: 14\n",
      "max lddt kept: 1.1537178652886946\n",
      "DPB1      176\n",
      "DQB1      153\n",
      "DRB1     7938\n",
      "DRB3      340\n",
      "DRB4      403\n",
      "DRB5      630\n",
      "IAB       292\n",
      "IEB        31\n",
      "other      23\n",
      "Name: locus, dtype: int64\n",
      "3048\n"
     ]
    }
   ],
   "source": [
    "#score distrib\n",
    "plt.figure()\n",
    "plt.hist(reduced_df['lddt_score'],histtype='step',bins=100,cumulative=1)\n",
    "plt.show()\n",
    "\n",
    "#reduce to top N\n",
    "N=10000\n",
    "df_top=reduced_df.sort_values('lddt_score').head(N).copy()\n",
    "df_top=df_top[df_top['register_identified']] #make sure register identified everywhere\n",
    "n=len(df_top)\n",
    "df_top=df_top[df_top['af_tails'].map(lambda x:(x[0]>=0)&(x[1]>=0))] #drop non-canonical tails\n",
    "print('non-canonical tails dropped:',n-len(df_top))\n",
    "print('max lddt kept:',df_top['lddt_score'].max())\n",
    "\n",
    "#look at locus distrib\n",
    "def map_locus(x):\n",
    "    y=x.info\n",
    "    if y['species'] in ['9606','10090']:\n",
    "        return y['locus']\n",
    "    else:\n",
    "        return 'other'\n",
    "df_top['locus']=df_top['mhc_b'].map(map_locus)\n",
    "print(df_top['locus'].value_counts().sort_index())\n",
    "\n",
    "#subsample DRB1\n",
    "df_top=pd.concat([df_top[df_top['locus']=='DRB1'].sample(1000),df_top[df_top['locus']!='DRB1']])\n",
    "print(len(df_top))\n",
    "\n",
    "#sort my MHC_B\n",
    "df_top['mhc_b_str']=df_top['mhc_b'].map(lambda x:'_'.join([x.info['species'],x.info['locus'],x.info['allele']]))\n",
    "df_top=df_top.sort_values('mhc_b_str')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "08408a11",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "62 Index(['class', 'cluster_seq', 'pdb_id', 'pmhc_id', 'date', 'tails',\n",
      "       'netmhc_tails', 'netmhc_kd', 'pep', 'mhc_a', 'mhc_b', 'exclude_pdbs',\n",
      "       'seqnn_logkds_all', 'seqnn_logkd', 'seqnn_tails', 'tails_prefiltered',\n",
      "       'templates'],\n",
      "      dtype='object')\n",
      "10     1\n",
      "11     2\n",
      "12     5\n",
      "13    15\n",
      "14    11\n",
      "15    12\n",
      "16     6\n",
      "18     1\n",
      "19     2\n",
      "20     3\n",
      "21     1\n",
      "22     2\n",
      "23     1\n",
      "Name: pep, dtype: int64\n",
      "[(9, 3) (2, 2) (4, 2) (3, 1) (3, 10) (2, 5) (1, 3) (3, 3) (2, 3) (8, 3)\n",
      " (7, 3) (3, 4) (3, 2) (0, 3) (4, 10) (-1, 2) (2, 1) (1, 2) (1, 1) (6, 5)\n",
      " (4, 5) (4, 3) (2, 0) (5, 1) (2, 4) (5, 2)]\n"
     ]
    }
   ],
   "source": [
    "#len and tail distrib in discovery\n",
    "target_df=pd.read_pickle('/scratch/vmikhayl/pmhc_modeling/discovery_dataset/run_10/target_df.pckl')\n",
    "target_df=target_df[target_df['class']=='II']\n",
    "print(len(target_df),target_df.columns)\n",
    "print(target_df['pep'].map(len).value_counts().sort_index())\n",
    "print(target_df['tails'].unique())\n",
    "\n",
    "#conclusion: take len from 10 to 23 (covers most of hist) and all tails\n",
    "#total pep registers:\n",
    "#2+3+..+15=17*7=119 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6e63afdf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total pep pdbnums: 119\n",
      "['   0v', '   0w', '   0x', '   0y', '   0z', '   01', '   02', '   03', '   04', '   05', '   06', '   07', '   08', '   09', '   1 ', '   2 ', '   3 ', '   4 ', '   5 ', '   6 ', '   7 ', '   8 ', '   9 ', '  10 ', '  11 ', '  12 ', '  13 ', '  14 ', '  15 ', '  16 ', '  17 ', '  18 ', '  19 ', '  20 ', '  21 ', '  22 ', '  23 ']\n",
      "pep pdbnums agree: True\n",
      "different cl II pdbnums:\n",
      "A 10\n",
      "B 44\n",
      "res positions in cl II pdbnums:\n",
      "A 85\n",
      "B 91\n",
      "{'A': array(['   02', '   01', '   1 ', '   2 ', '   3 ', '   4 ', '   5 ',\n",
      "       '   6 ', '   7 ', '   7A', '   8 ', '   9 ', '  10 ', '  11 ',\n",
      "       '  12 ', '  13 ', '  14 ', '  19 ', '  20 ', '  21 ', '  22 ',\n",
      "       '  23 ', '  24 ', '  25 ', '  26 ', '  27 ', '  28 ', '  29 ',\n",
      "       '  30 ', '  31 ', '  32 ', '  33 ', '  34 ', '  35 ', '  36 ',\n",
      "       '  37 ', '  38 ', '  39 ', '  42 ', '  43 ', '  44 ', '  45 ',\n",
      "       '  46 ', '  47 ', '  48 ', '  49 ', '  50 ', '  51 ', '  52 ',\n",
      "       '  53 ', '  54 ', '  59 ', '  60 ', '  61 ', '  62 ', '  63 ',\n",
      "       '  64 ', '  65 ', '  66 ', '  67 ', '  68 ', '  69 ', '  70 ',\n",
      "       '  71 ', '  72 ', '  73 ', '  74 ', '  75 ', '  76 ', '  77 ',\n",
      "       '  78 ', '  79 ', '  80 ', '  81 ', '  82 ', '  83 ', '  84 ',\n",
      "       '  85 ', '  86 ', '  87 ', '  88 ', '  89 ', '  90 ', '  91 ',\n",
      "       '  92 '], dtype='<U5'), 'B': array(['1001 ', '1002 ', '1003 ', '1004 ', '1005 ', '1006 ', '1007 ',\n",
      "       '1008 ', '1009 ', '1010 ', '1011 ', '1012 ', '1013 ', '1014 ',\n",
      "       '1015 ', '1016 ', '1017 ', '1018 ', '1019 ', '1020 ', '1021 ',\n",
      "       '1022 ', '1023 ', '1024 ', '1025 ', '1026 ', '1027 ', '1028 ',\n",
      "       '1029 ', '1030 ', '1031 ', '1032 ', '1033 ', '1034 ', '1035 ',\n",
      "       '1036 ', '1037 ', '1038 ', '1039 ', '1042 ', '1043 ', '1044 ',\n",
      "       '1045 ', '1046 ', '1047 ', '1048 ', '1049 ', '1050 ', '1051 ',\n",
      "       '1052 ', '1053 ', '1054 ', '1055 ', '1056 ', '1057 ', '1058 ',\n",
      "       '1059 ', '1060 ', '1061 ', '1061A', '1061B', '1062 ', '1063 ',\n",
      "       '1064 ', '1065 ', '1066 ', '1067 ', '1068 ', '1069 ', '1070 ',\n",
      "       '1071 ', '1072 ', '1072A', '1073 ', '1074 ', '1075 ', '1076 ',\n",
      "       '1077 ', '1078 ', '1079 ', '1080 ', '1081 ', '1082 ', '1083 ',\n",
      "       '1084 ', '1085 ', '1086 ', '1087 ', '1088 ', '1089 ', '1090 '],\n",
      "      dtype='<U5')}\n"
     ]
    }
   ],
   "source": [
    "#make master pdbnums for mhc and pep\n",
    "\n",
    "###pep###\n",
    "import tfold.nn.nn_utils as nn_utils\n",
    "import tfold.modeling.template_tools as template_tools\n",
    "\n",
    "lmin,lmax=10,23 #choose min and max pep length; all >=0 tails will be considered\n",
    "\n",
    "pep_pdbnum=set()\n",
    "pep_pdbnums_all=set()\n",
    "for l in range(lmin,lmax+1):\n",
    "    for l_tail in range(l-8):\n",
    "        c_pdbnum=template_tools._make_pep_pdbnums_II(l,l_tail)\n",
    "        pep_pdbnums_all.add(tuple(c_pdbnum))\n",
    "        pep_pdbnum|=set(c_pdbnum)\n",
    "pep_pdbnums_all=[list(x) for x in pep_pdbnums_all]\n",
    "pep_pdbnum0=sorted(list(pep_pdbnum))\n",
    "pep_pdbnum=template_tools._make_pep_pdbnums_II(2*lmax-9,lmax-9) #this way it is sorted\n",
    "print('total pep pdbnums:',len(pep_pdbnums_all))\n",
    "print(pep_pdbnum)\n",
    "print('pep pdbnums agree:',set(pep_pdbnum0)==set(pep_pdbnum))\n",
    "\n",
    "###mhc###\n",
    "#NOTE: in MHC MSAs, used only pdbnums for mhcs in our structure database; here using all\n",
    "\n",
    "#collect all cl II mhc pdbnums\n",
    "mhc_pdbnums_all={'A':set(),'B':set()} #by chain\n",
    "for mhc in seq_tools.mhcs.values():\n",
    "    if mhc.info['class']=='II':\n",
    "        mhc_pdbnums_all[mhc.info['chain']].add(tuple(mhc.data['pdbnum']))\n",
    "for k in mhc_pdbnums_all:\n",
    "    mhc_pdbnums_all[k]=[list(x) for x in mhc_pdbnums_all[k]]\n",
    "print('different cl II pdbnums:')\n",
    "for k in mhc_pdbnums_all:\n",
    "    print(k,len(mhc_pdbnums_all[k]))\n",
    "\n",
    "#all res positions in them\n",
    "pdbnum_entries={'A':set(),'B':set()}\n",
    "for k in pdbnum_entries:\n",
    "    for pdbnum in mhc_pdbnums_all[k]:\n",
    "        pdbnum_entries[k]|=set(pdbnum)\n",
    "    pdbnum_entries[k]=np.sort(list(pdbnum_entries[k]))\n",
    "print('res positions in cl II pdbnums:')\n",
    "for k in pdbnum_entries:\n",
    "    print(k,len(pdbnum_entries[k]))\n",
    "\n",
    "#permute to ensure order\n",
    "for k in mhc_pdbnums_all:\n",
    "    for pdbnum in mhc_pdbnums_all[k]:\n",
    "        pdbnum_entries[k][np.isin(pdbnum_entries[k],pdbnum)]=pdbnum\n",
    "#check\n",
    "for pdbnum in mhc_pdbnums_all[k]:\n",
    "    ind=np.isin(pdbnum_entries[k],pdbnum)\n",
    "    if np.sum(ind)!=len(pdbnum):\n",
    "        print('oops, missing')\n",
    "    if '|'.join(pdbnum_entries[k][ind])!='|'.join(pdbnum):\n",
    "        print('oops, order')\n",
    "mhc_pdbnum={}\n",
    "for k in pdbnum_entries: \n",
    "    mhc_pdbnum[k]=np.array(pdbnum_entries[k])\n",
    "print(mhc_pdbnum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "eb4507e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3048, 37) (3048, 85) (3048, 91)\n"
     ]
    }
   ],
   "source": [
    "#convert peps and mhcs in assays to NUMSEQ, extend to master pdbnum\n",
    "def pep_to_numseq(x):\n",
    "    seq=x['pep']\n",
    "    tails=x['af_tails']\n",
    "    return seq_tools.NUMSEQ(seq=seq,pdbnum=template_tools._make_pep_pdbnums_II(len(seq),tails[0]))\n",
    "def extend_numseq(s,pdbnum):\n",
    "    s_ext=seq_tools.NUMSEQ(seq='_'*len(pdbnum),pdbnum=pdbnum) #use '_' for gap symbol, to replace later\n",
    "    ind1=np.isin(pdbnum,s.data['pdbnum'])\n",
    "    ind2=np.isin(s.data['pdbnum'],pdbnum)\n",
    "    s_ext.data[ind1]=s.data[ind2]\n",
    "    return s_ext.data['seq']\n",
    "pep_assays=df_top.apply(pep_to_numseq,axis=1)\n",
    "pep_assays=np.stack(pep_assays.map(lambda x: extend_numseq(x,pep_pdbnum)).values)\n",
    "mhc_assays_a=df_top['mhc_a'].copy()\n",
    "mhc_assays_a=np.stack(mhc_assays_a.map(lambda x: extend_numseq(x,mhc_pdbnum['A'])).values)\n",
    "mhc_assays_b=df_top['mhc_b'].copy()\n",
    "mhc_assays_b=np.stack(mhc_assays_b.map(lambda x: extend_numseq(x,mhc_pdbnum['B'])).values)\n",
    "print(pep_assays.shape,mhc_assays_a.shape,mhc_assays_b.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "dc020ea4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "completed in 16.224181652069092\n"
     ]
    }
   ],
   "source": [
    "#reduce to a3m, save pdbnum->filename index\n",
    "\n",
    "pmhc_msa_index={'II_P':{},'II_M':{},'II_N':{}} #cl_chain: pdbnum: id        \n",
    "\n",
    "lower_vect=np.vectorize(lambda s:s.lower())\n",
    "repl_gap_vect=np.vectorize(lambda s:s.replace('_','-'))\n",
    "def reduce_matrix_to_pdbnum(matrix,ind):\n",
    "    '''\n",
    "    entries within pdbnum are kept, with gap symbols replaced by '-';\n",
    "    entries outside pdbnum are made lowercase, with gap symbols dropped;\n",
    "    outputs an array of strings\n",
    "    '''\n",
    "    matrix=matrix.copy()\n",
    "    if np.any(~ind):\n",
    "        matrix[:,~ind]=lower_vect(matrix[:,~ind])\n",
    "    if np.any(ind):\n",
    "        matrix[:,ind]=repl_gap_vect(matrix[:,ind])\n",
    "    return matrix\n",
    "\n",
    "t0=time.time()\n",
    "#pep\n",
    "for i,pdbnum in enumerate(pep_pdbnums_all):\n",
    "    ind=np.isin(pep_pdbnum,pdbnum) #which to keep\n",
    "    pep_assays_reduced=reduce_matrix_to_pdbnum(pep_assays,ind)\n",
    "    pep_a3m=[]\n",
    "    for j,line in enumerate(pep_assays_reduced):\n",
    "        pep_a3m.append(f'>{j}')\n",
    "        pep_a3m.append(''.join(line).replace('_',''))        \n",
    "    with open(f'./data/msas/pMHC/II_P_{i}.a3m','w') as f:\n",
    "        f.write('\\n'.join(pep_a3m))\n",
    "    pmhc_msa_index['II_P']['|'.join(pdbnum)]=i\n",
    "    \n",
    "#mhc\n",
    "mhc_assays={'A':mhc_assays_a,'B':mhc_assays_b}\n",
    "name_dict={'A':'M','B':'N'}\n",
    "for k in ['A','B']:\n",
    "    for i,pdbnum in enumerate(mhc_pdbnums_all[k]):\n",
    "        ind=np.isin(mhc_pdbnum[k],pdbnum) #which to keep\n",
    "        mhc_assays_reduced=reduce_matrix_to_pdbnum(mhc_assays[k],ind)\n",
    "        mhc_a3m=[]\n",
    "        for j,line in enumerate(mhc_assays_reduced):\n",
    "            mhc_a3m.append(f'>{j}')\n",
    "            mhc_a3m.append(''.join(line).replace('_',''))        \n",
    "        with open(f'./data/msas/pMHC/II_{name_dict[k]}_{i}.a3m','w') as f:\n",
    "            f.write('\\n'.join(mhc_a3m))\n",
    "        pmhc_msa_index[f'II_{name_dict[k]}']['|'.join(pdbnum)]=i\n",
    "\n",
    "#update lookup table\n",
    "with open('./data/msas/pmhc_msa_index.pckl','rb') as f:\n",
    "    pmhc_msa_index_old=pickle.load(f)\n",
    "pmhc_msa_index.update(pmhc_msa_index_old)\n",
    "with open('./data/msas/pmhc_msa_index.pckl','wb') as f:\n",
    "    pickle.dump(pmhc_msa_index,f)\n",
    "print('completed in {}'.format(time.time()-t0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1187487e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}