20230202 preprocess_try.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7174a1f8-6db5-4e26-83f7-ef4390d10050",
   "metadata": {},
   "source": [
    "# Seq PDB map try"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "80f69057-0719-4472-a418-3fa15f77b622",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df_davis = pd.read_csv('./davis_seq_pdb.csv')\n",
    "df_kiba = pd.read_csv('./kiba_seq_pdb.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a8a2ad02-0967-4549-bb85-b79a9ea8de07",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'QLTPTNSLKRGGAHHRRCEVALLGCGAVLAATGLGFDLLEAGKCQLLPLEEPEPPAREEKKRREGLFQRSSRPRRSTSPPSRKLFKKEEPMLLLGDPSASLTLLSLSSISECNSTRSLLRSDSDEIVVYEMPVSPVEAPPLSPCTHNPLVNVRVERFKRDPNQSLTPTHVTLTTPSQPSSHRRTPSDGALKPETLLASRSPSSNGLSPSPGAGMLKTPSPSRDPGEFPRLPDPNVVFPPTPRRWNTQQDSTLERPKTLEFLPRPRPSANRQRLDPWWFVSPSHARSTSPANSSSTETPSNLDSCFASSSSTVEERPGLPALLPFQAGPLPPTERTLLDLDAEGQSQDSTVPLCRAELNTHRPAPYEIQQEFWS'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_davis[df_davis['target_pdb_id'] == 'TOFOLD']['target_sequence'].to_list()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "50254a09-c9c4-40f7-baea-7306bb634c87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6G5I_36\n"
     ]
    }
   ],
   "source": [
    "for pdb_id in df_davis['target_pdb_id'].unique():\n",
    "    if pdb_id[: 4].lower() in [\"6g5i\", \"5t0j\", \"5wve\"]:\n",
    "        print(pdb_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "689b9f6b-6214-4c3a-8e4b-3e8e9bc07b71",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pdb_id in df_kiba['target_pdb_id'].unique():\n",
    "    if pdb_id[: 4].lower() in [\"6g5i\", \"5t0j\", \"5wve\"]:\n",
    "        print(pdb_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fc3ead27-e2f4-4630-ac29-111a3bf24cf1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "298"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df_davis['target_pdb_id'].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "acd9c456-0889-4fde-a5bf-59144b89c141",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "298"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set(pdb_id[: 4].lower() for pdb_id in df_davis['target_pdb_id'].unique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e5f184bf-910e-4ac1-ace2-a970baf9a4ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "\n",
    "for fn in glob.glob('./data/**/*.csv', recursive=True):\n",
    "    df = pd.read_csv(fn)\n",
    "    df = df.rename(columns={col: col.strip() for col in df.columns})\n",
    "    df.to_csv(fn, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d52a062-a165-451c-903e-327ffcf4d570",
   "metadata": {},
   "source": [
    "# Prot. -> LLM (FusionDTA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d7d6f3a7-6388-4aeb-be4a-31728674b72a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import json,pickle\n",
    "from collections import OrderedDict\n",
    "import re\n",
    "import csv\n",
    "import esm\n",
    "import torch\n",
    "from rdkit import Chem\n",
    "import pypdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0f7744a8-2f97-43e3-8966-35fde2a3597a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_protein_pretraining_representation(dataset_name, prots):\n",
    "    data_dict = {}\n",
    "    prots_tuple = [(str(i), prots[i][:1022]) for i in range(len(prots))]\n",
    "    model, alphabet = esm.pretrained.esm2_t36_3B_UR50D()\n",
    "    batch_converter = alphabet.get_batch_converter()\n",
    "    i = 0\n",
    "    batch = 1\n",
    "    \n",
    "    while (batch*i) < len(prots):\n",
    "        print('converting protein batch: '+ str(i))\n",
    "        if (i + batch) < len(prots):\n",
    "            pt = prots_tuple[batch*i:batch*(i+1)]\n",
    "        else:\n",
    "            pt = prots_tuple[batch*i:]\n",
    "        \n",
    "        batch_labels, batch_strs, batch_tokens = batch_converter(pt)\n",
    "        #model = model.cuda()\n",
    "        #batch_tokens = batch_tokens.cuda()\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            results = model(batch_tokens, repr_layers=[36], return_contacts=True)\n",
    "        token_representations = results[\"representations\"][36].numpy()\n",
    "        data_dict[i] = token_representations\n",
    "        i += 1\n",
    "    np.savez(dataset_name + '.npz', dict=data_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "51fd710b-c68e-4ad0-a2a4-28e8dbe9566f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seq = 'MSPLNQSAEGLPQEASNRSLNATETSEAWDPRTLQALKISLAVVLSVITLATVLSNAFVLTTILLTRKLHTPANYLIGSLATTDLLVSILVMPISIAYTITHTWNFGQILCDIWLSSDITCCTASILHLCVIALDRYWAITDALEYSKRRTAGHAATMIAIVWAISICISIPPLFWRQAKAQEEMSDCLVNTSQISYTIYSTCGAFYIPSVLLIILYGRIYRAARNRILNPPSLYGKRFTTAHLITGSAGSSLCSLNSSLHEGHSHSAGSPLFFNHVKIKLADSALERKRISAARERKATKILGIILGAFIICWLPFFVVSLVLPICRDSCWIHPALFDFFTWLGYLNSLINPIIYTVFNEEFRQAFQKIVPFRKAS'\n",
    "q = pypdb.Query(seq,\n",
    "          query_type=\"sequence\", \n",
    "          return_type=\"polymer_entity\")\n",
    "None is None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "174b6fa3-b1d7-4f0d-aafd-82a02a5b47b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_pdb_id(seq):\n",
    "    q = pypdb.Query(seq,\n",
    "              query_type=\"sequence\", \n",
    "              return_type=\"polymer_entity\")\n",
    "    res = q.search()\n",
    "    \n",
    "    if res is not None:\n",
    "        return res['result_set'][0]['identifier']\n",
    "    else:\n",
    "        return 'TOFOLD'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c0f4ea91-8240-4bde-ab50-7774bb070d1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 68/68 [00:00<00:00, 2862.29it/s]\n",
      " 57%|███████████████████████████████████████████████████████████████▌                                               | 253/442 [15:59<11:22,  3.61s/it]/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/pypdb/util/http_requests.py:65: UserWarning: Too many failures on requests. Exiting...\n",
      "  warnings.warn(\"Too many failures on requests. Exiting...\")\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/pypdb/pypdb.py:293: UserWarning: Retrieval failed, returning None\n",
      "  warnings.warn(\"Retrieval failed, returning None\")\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 442/442 [26:31<00:00,  3.60s/it]\n"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'fused_layer_norm_cuda'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 39\u001b[0m\n\u001b[1;32m     37\u001b[0m \u001b[38;5;66;03m# protein pretraing presentation\u001b[39;00m\n\u001b[1;32m     38\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m ESM:\n\u001b[0;32m---> 39\u001b[0m     \u001b[43mgenerate_protein_pretraining_representation\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprot_seqs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     41\u001b[0m affinity \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39masarray(affinity)\n\u001b[1;32m     42\u001b[0m opts \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m,\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvalid\u001b[39m\u001b[38;5;124m'\u001b[39m]\n",
      "Cell \u001b[0;32mIn[2], line 4\u001b[0m, in \u001b[0;36mgenerate_protein_pretraining_representation\u001b[0;34m(dataset_name, prots)\u001b[0m\n\u001b[1;32m      2\u001b[0m data_dict \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m      3\u001b[0m prots_tuple \u001b[38;5;241m=\u001b[39m [(\u001b[38;5;28mstr\u001b[39m(i), prots[i][:\u001b[38;5;241m1022\u001b[39m]) \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(prots))]\n\u001b[0;32m----> 4\u001b[0m model, alphabet \u001b[38;5;241m=\u001b[39m \u001b[43mesm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpretrained\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mesm2_t36_3B_UR50D\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      5\u001b[0m batch_converter \u001b[38;5;241m=\u001b[39m alphabet\u001b[38;5;241m.\u001b[39mget_batch_converter()\n\u001b[1;32m      6\u001b[0m i \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/pretrained.py:387\u001b[0m, in \u001b[0;36mesm2_t36_3B_UR50D\u001b[0;34m()\u001b[0m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mesm2_t36_3B_UR50D\u001b[39m():\n\u001b[1;32m    383\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"36 layer ESM-2 model with 3B params, trained on UniRef50.\u001b[39;00m\n\u001b[1;32m    384\u001b[0m \n\u001b[1;32m    385\u001b[0m \u001b[38;5;124;03m    Returns a tuple of (Model, Alphabet).\u001b[39;00m\n\u001b[1;32m    386\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 387\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mload_model_and_alphabet_hub\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mesm2_t36_3B_UR50D\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/pretrained.py:64\u001b[0m, in \u001b[0;36mload_model_and_alphabet_hub\u001b[0;34m(model_name)\u001b[0m\n\u001b[1;32m     62\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mload_model_and_alphabet_hub\u001b[39m(model_name):\n\u001b[1;32m     63\u001b[0m     model_data, regression_data \u001b[38;5;241m=\u001b[39m _download_model_and_regression_data(model_name)\n\u001b[0;32m---> 64\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mload_model_and_alphabet_core\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel_data\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mregression_data\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/pretrained.py:191\u001b[0m, in \u001b[0;36mload_model_and_alphabet_core\u001b[0;34m(model_name, model_data, regression_data)\u001b[0m\n\u001b[1;32m    188\u001b[0m     model_data[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mupdate(regression_data[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[1;32m    190\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m model_name\u001b[38;5;241m.\u001b[39mstartswith(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mesm2\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 191\u001b[0m     model, alphabet, model_state \u001b[38;5;241m=\u001b[39m \u001b[43m_load_model_and_alphabet_core_v2\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel_data\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    192\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    193\u001b[0m     model, alphabet, model_state \u001b[38;5;241m=\u001b[39m _load_model_and_alphabet_core_v1(model_data)\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/pretrained.py:176\u001b[0m, in \u001b[0;36m_load_model_and_alphabet_core_v2\u001b[0;34m(model_data)\u001b[0m\n\u001b[1;32m    174\u001b[0m state_dict \u001b[38;5;241m=\u001b[39m upgrade_state_dict(state_dict)\n\u001b[1;32m    175\u001b[0m alphabet \u001b[38;5;241m=\u001b[39m esm\u001b[38;5;241m.\u001b[39mdata\u001b[38;5;241m.\u001b[39mAlphabet\u001b[38;5;241m.\u001b[39mfrom_architecture(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mESM-1b\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 176\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mESM2\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    177\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnum_layers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcfg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoder_layers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    178\u001b[0m \u001b[43m    \u001b[49m\u001b[43membed_dim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcfg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoder_embed_dim\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    179\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattention_heads\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcfg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoder_attention_heads\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    180\u001b[0m \u001b[43m    \u001b[49m\u001b[43malphabet\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43malphabet\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    181\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtoken_dropout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcfg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtoken_dropout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    182\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    183\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m model, alphabet, state_dict\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/model/esm2.py:39\u001b[0m, in \u001b[0;36mESM2.__init__\u001b[0;34m(self, num_layers, embed_dim, attention_heads, alphabet, token_dropout)\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mappend_eos \u001b[38;5;241m=\u001b[39m alphabet\u001b[38;5;241m.\u001b[39mappend_eos\n\u001b[1;32m     37\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtoken_dropout \u001b[38;5;241m=\u001b[39m token_dropout\n\u001b[0;32m---> 39\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_init_submodules\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/model/esm2.py:50\u001b[0m, in \u001b[0;36mESM2._init_submodules\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_scale \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m     43\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_tokens \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mEmbedding(\n\u001b[1;32m     44\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39malphabet_size,\n\u001b[1;32m     45\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim,\n\u001b[1;32m     46\u001b[0m     padding_idx\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpadding_idx,\n\u001b[1;32m     47\u001b[0m )\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlayers \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mModuleList(\n\u001b[0;32m---> 50\u001b[0m     [\n\u001b[1;32m     51\u001b[0m         TransformerLayer(\n\u001b[1;32m     52\u001b[0m             \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim,\n\u001b[1;32m     53\u001b[0m             \u001b[38;5;241m4\u001b[39m \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim,\n\u001b[1;32m     54\u001b[0m             \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattention_heads,\n\u001b[1;32m     55\u001b[0m             add_bias_kv\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     56\u001b[0m             use_esm1b_layer_norm\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     57\u001b[0m             use_rotary_embeddings\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     58\u001b[0m         )\n\u001b[1;32m     59\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_layers)\n\u001b[1;32m     60\u001b[0m     ]\n\u001b[1;32m     61\u001b[0m )\n\u001b[1;32m     63\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontact_head \u001b[38;5;241m=\u001b[39m ContactPredictionHead(\n\u001b[1;32m     64\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_layers \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattention_heads,\n\u001b[1;32m     65\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprepend_bos,\n\u001b[1;32m     66\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mappend_eos,\n\u001b[1;32m     67\u001b[0m     eos_idx\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39meos_idx,\n\u001b[1;32m     68\u001b[0m )\n\u001b[1;32m     69\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39memb_layer_norm_after \u001b[38;5;241m=\u001b[39m ESM1bLayerNorm(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim)\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/model/esm2.py:51\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_scale \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m     43\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_tokens \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mEmbedding(\n\u001b[1;32m     44\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39malphabet_size,\n\u001b[1;32m     45\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim,\n\u001b[1;32m     46\u001b[0m     padding_idx\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpadding_idx,\n\u001b[1;32m     47\u001b[0m )\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlayers \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mModuleList(\n\u001b[1;32m     50\u001b[0m     [\n\u001b[0;32m---> 51\u001b[0m         \u001b[43mTransformerLayer\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     52\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43membed_dim\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     53\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43membed_dim\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     54\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattention_heads\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     55\u001b[0m \u001b[43m            \u001b[49m\u001b[43madd_bias_kv\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     56\u001b[0m \u001b[43m            \u001b[49m\u001b[43muse_esm1b_layer_norm\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     57\u001b[0m \u001b[43m            \u001b[49m\u001b[43muse_rotary_embeddings\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     58\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     59\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_layers)\n\u001b[1;32m     60\u001b[0m     ]\n\u001b[1;32m     61\u001b[0m )\n\u001b[1;32m     63\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontact_head \u001b[38;5;241m=\u001b[39m ContactPredictionHead(\n\u001b[1;32m     64\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_layers \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattention_heads,\n\u001b[1;32m     65\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprepend_bos,\n\u001b[1;32m     66\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mappend_eos,\n\u001b[1;32m     67\u001b[0m     eos_idx\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39meos_idx,\n\u001b[1;32m     68\u001b[0m )\n\u001b[1;32m     69\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39memb_layer_norm_after \u001b[38;5;241m=\u001b[39m ESM1bLayerNorm(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim)\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/modules.py:101\u001b[0m, in \u001b[0;36mTransformerLayer.__init__\u001b[0;34m(self, embed_dim, ffn_embed_dim, attention_heads, add_bias_kv, use_esm1b_layer_norm, use_rotary_embeddings)\u001b[0m\n\u001b[1;32m     99\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattention_heads \u001b[38;5;241m=\u001b[39m attention_heads\n\u001b[1;32m    100\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muse_rotary_embeddings \u001b[38;5;241m=\u001b[39m use_rotary_embeddings\n\u001b[0;32m--> 101\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_init_submodules\u001b[49m\u001b[43m(\u001b[49m\u001b[43madd_bias_kv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43muse_esm1b_layer_norm\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/esm/modules.py:113\u001b[0m, in \u001b[0;36mTransformerLayer._init_submodules\u001b[0;34m(self, add_bias_kv, use_esm1b_layer_norm)\u001b[0m\n\u001b[1;32m    104\u001b[0m BertLayerNorm \u001b[38;5;241m=\u001b[39m ESM1bLayerNorm \u001b[38;5;28;01mif\u001b[39;00m use_esm1b_layer_norm \u001b[38;5;28;01melse\u001b[39;00m ESM1LayerNorm\n\u001b[1;32m    106\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mself_attn \u001b[38;5;241m=\u001b[39m MultiheadAttention(\n\u001b[1;32m    107\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim,\n\u001b[1;32m    108\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattention_heads,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    111\u001b[0m     use_rotary_embeddings\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muse_rotary_embeddings,\n\u001b[1;32m    112\u001b[0m )\n\u001b[0;32m--> 113\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mself_attn_layer_norm \u001b[38;5;241m=\u001b[39m \u001b[43mBertLayerNorm\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43membed_dim\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    115\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfc1 \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mLinear(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mffn_embed_dim)\n\u001b[1;32m    116\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfc2 \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mLinear(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mffn_embed_dim, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim)\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/site-packages/apex-0.1-py3.10-linux-x86_64.egg/apex/normalization/fused_layer_norm.py:268\u001b[0m, in \u001b[0;36mFusedLayerNorm.__init__\u001b[0;34m(self, normalized_shape, eps, elementwise_affine)\u001b[0m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m()\n\u001b[1;32m    267\u001b[0m \u001b[38;5;28;01mglobal\u001b[39;00m fused_layer_norm_cuda\n\u001b[0;32m--> 268\u001b[0m fused_layer_norm_cuda \u001b[38;5;241m=\u001b[39m \u001b[43mimportlib\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimport_module\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfused_layer_norm_cuda\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    270\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(normalized_shape, numbers\u001b[38;5;241m.\u001b[39mIntegral):\n\u001b[1;32m    271\u001b[0m     normalized_shape \u001b[38;5;241m=\u001b[39m (normalized_shape,)\n",
      "File \u001b[0;32m~/miniconda3/envs/pl102/lib/python3.10/importlib/__init__.py:126\u001b[0m, in \u001b[0;36mimport_module\u001b[0;34m(name, package)\u001b[0m\n\u001b[1;32m    124\u001b[0m             \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m    125\u001b[0m         level \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[0;32m--> 126\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_bootstrap\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_gcd_import\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m[\u001b[49m\u001b[43mlevel\u001b[49m\u001b[43m:\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpackage\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlevel\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m<frozen importlib._bootstrap>:1050\u001b[0m, in \u001b[0;36m_gcd_import\u001b[0;34m(name, package, level)\u001b[0m\n",
      "File \u001b[0;32m<frozen importlib._bootstrap>:1027\u001b[0m, in \u001b[0;36m_find_and_load\u001b[0;34m(name, import_)\u001b[0m\n",
      "File \u001b[0;32m<frozen importlib._bootstrap>:1004\u001b[0m, in \u001b[0;36m_find_and_load_unlocked\u001b[0;34m(name, import_)\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'fused_layer_norm_cuda'"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "ESM = True\n",
    "datasets = ['davis','kiba']\n",
    "for dataset in datasets:\n",
    "    fpath = 'data/' + dataset + '/raw/'\n",
    "    train_valid_folds = json.load(open(fpath + \"folds/train_fold_setting1.txt\"))\n",
    "    test_fold = json.load(open(fpath + \"folds/test_fold_setting1.txt\"))\n",
    "    valid_ids = [0, 1, 2, 3, 4]\n",
    "    valid_folds = [train_valid_folds[vid] for vid in valid_ids]\n",
    "    train_folds = []\n",
    "    for valid_id in valid_ids:\n",
    "        temp = []\n",
    "        for idx in range(5):\n",
    "            if idx != valid_id:\n",
    "                temp += train_valid_folds[idx]\n",
    "        train_folds.append(temp)\n",
    "    \n",
    "    ligands = json.load(open(fpath + \"ligands_can.txt\"), object_pairs_hook=OrderedDict)\n",
    "    proteins = json.load(open(fpath + \"proteins.txt\"), object_pairs_hook=OrderedDict)\n",
    "    affinity = pickle.load(open(fpath + \"Y\",\"rb\"), encoding='latin1')\n",
    "    drugs = []\n",
    "    drug_smiles = []\n",
    "    prot_seqs = []\n",
    "    prot_pdb_ids = []\n",
    "    for d in tqdm(ligands.keys()):\n",
    "        #lg = ligands[d]\n",
    "        lg = Chem.MolToSmiles(Chem.MolFromSmiles(ligands[d]),isomericSmiles=True)\n",
    "        drugs.append(lg)\n",
    "        drug_smiles.append(ligands[d])\n",
    "    for t in tqdm(proteins.keys()):\n",
    "        prot_seqs.append(proteins[t])\n",
    "        prot_pdb_ids.append(query_pdb_id(proteins[t]))\n",
    "    if dataset == 'davis':\n",
    "        affinity = [-np.log10(y/1e9) for y in affinity]\n",
    "\n",
    "    # protein pretraing presentation\n",
    "    if ESM:\n",
    "        generate_protein_pretraining_representation(dataset, prot_seqs)\n",
    "\n",
    "    affinity = np.asarray(affinity)\n",
    "    opts = ['train','valid']\n",
    "\n",
    "    print('generating test data')\n",
    "    rows, cols = np.where(np.isnan(affinity) == False)  \n",
    "    test_rows, test_cols = rows[test_fold], cols[test_fold]\n",
    "    with open('data/' + dataset + '/' + dataset  + '_test.csv', 'w') as f:\n",
    "        f.write('compound_iso_smiles, target_sequence, target_pdb_id, affinity, protein_id, drug_id\\n')\n",
    "        for pair_ind in range(len(test_rows)):\n",
    "            ls = []\n",
    "            ls += [ drugs[test_rows[pair_ind]] ]\n",
    "            ls += [ prot_seqs[test_cols[pair_ind]] ]\n",
    "            ls += [ prot_pdb_ids[test_cols[pair_ind]] ]\n",
    "            ls += [ affinity[test_rows[pair_ind], test_cols[pair_ind]] ]\n",
    "            ls += [ test_cols[pair_ind] ]\n",
    "            ls += [ test_rows[pair_ind] ]\n",
    "            f.write(','.join(map(str, ls)) + '\\n')\n",
    "\n",
    "    for i in range(5):\n",
    "        train_fold = train_folds[i]\n",
    "        valid_fold = valid_folds[i]\n",
    "        for opt in opts:\n",
    "            rows, cols = np.where(np.isnan(affinity) == False)  \n",
    "            if opt == 'train':\n",
    "                rows, cols = rows[train_fold], cols[train_fold]\n",
    "                \n",
    "                #generating cold data\n",
    "                with open('data/' + dataset + '_cold' + '.csv', 'w') as f:\n",
    "                    f.write('compound_iso_smiles, target_sequence, affinity, protein_id, drug_id\\n')\n",
    "                    for pair_ind in range(len(rows)):\n",
    "                        ls = []\n",
    "                        ls += [ drugs[rows[pair_ind]]  ]\n",
    "                        ls += [ prots[cols[pair_ind]]  ]\n",
    "                        ls += [ affinity[rows[pair_ind],cols[pair_ind]]  ]\n",
    "                        ls += [ cols[pair_ind] ]\n",
    "                        ls += [ rows[pair_ind] ]\n",
    "                        f.write(','.join(map(str,ls)) + '\\n') \n",
    "            elif opt == 'valid':\n",
    "                rows, cols = rows[valid_fold], cols[valid_fold]\n",
    "                \n",
    "                #generating cold data\n",
    "                with open('data/' + dataset + '_cold' + '.csv', 'a') as f:\n",
    "                    for pair_ind in range(len(rows)):\n",
    "                        ls = []\n",
    "                        ls += [ drugs[rows[pair_ind]]  ]\n",
    "                        ls += [ prots[cols[pair_ind]]  ]\n",
    "                        ls += [ affinity[rows[pair_ind],cols[pair_ind]]  ]\n",
    "                        ls += [ cols[pair_ind] ]\n",
    "                        ls += [ rows[pair_ind] ]\n",
    "                        f.write(','.join(map(str,ls)) + '\\n') \n",
    "                      \n",
    "            #5-fold data\n",
    "            print('generating 5-fold data')\n",
    "            with open('data/' + dataset + '/' + dataset + '_' + opt + '_fold_' + str(i) + '.csv', 'w') as f:\n",
    "                f.write('compound_iso_smiles, target_sequence, target_pdb_id, affinity, protein_id, drug_id\\n')\n",
    "                for pair_ind in range(len(rows)):\n",
    "                    ls = []\n",
    "                    ls += [ drugs[rows[pair_ind]] ]\n",
    "                    ls += [ prot_seqs[cols[pair_ind]] ]\n",
    "                    ls += [ prot_pdb_ids[cols[pair_ind]] ]\n",
    "                    ls += [ affinity[rows[pair_ind], cols[pair_ind]] ]\n",
    "                    ls += [ cols[pair_ind] ]\n",
    "                    ls += [ rows[pair_ind] ]\n",
    "                    f.write(','.join(map(str,ls)) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f32eec-1677-4eb8-9c6b-9f406377b76f",
   "metadata": {},
   "source": [
    "# Prot. -> PDB file -> Hybrid Graphs (AttnSiteDTI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "53466332-c7e5-4e43-ac95-61c43cb263bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-02-02 15:04:47.591148: 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",
      "2023-02-02 15:04:48.298870: 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",
      "2023-02-02 15:04:52.212567: 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; LD_LIBRARY_PATH: /home/luozc/CCMpred/lib:/home/luozc/miniconda3/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:\n",
      "2023-02-02 15:04:52.212859: 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; LD_LIBRARY_PATH: /home/luozc/CCMpred/lib:/home/luozc/miniconda3/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:\n",
      "2023-02-02 15:04:52.212872: 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"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pickle\n",
    "from collections import OrderedDict\n",
    "import random\n",
    "import glob\n",
    "\n",
    "import pandas as pd\n",
    "from dgllife.utils import smiles_to_bigraph, CanonicalAtomFeaturizer\n",
    "import dgl\n",
    "import numpy as np\n",
    "import torch\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem.rdmolops import GetAdjacencyMatrix\n",
    "import networkx as nx\n",
    "from Bio.PDB import *\n",
    "import deepchem\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f36df39-d358-4c17-9da3-f87a31c1c3d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_of_k_encoding(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        raise Exception(\"input {0} not in allowable set{1}:\".format(x, allowable_set))\n",
    "    # print list((map(lambda s: x == s, allowable_set)))\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def one_of_k_encoding_unk(x, allowable_set):\n",
    "    \"\"\"Maps inputs not in the allowable set to the last element.\"\"\"\n",
    "    if x not in allowable_set:\n",
    "        x = allowable_set[-1]\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def atom_feature(atom):\n",
    "    return np.array(one_of_k_encoding_unk(atom.GetSymbol(),\n",
    "                                          ['C', 'N', 'O', 'S', 'F', 'P', 'Cl', 'Br', 'B', 'H']) +\n",
    "                    one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5, 6, 7, 8]) +\n",
    "                    one_of_k_encoding_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4]) +\n",
    "                    one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5]) +\n",
    "                    [atom.GetIsAromatic()])  # (10, 9, 5, 6, 1) --> total 31\n",
    "\n",
    "\n",
    "def get_atom_feature(m):\n",
    "    H = []\n",
    "    for i in range(len(m)):\n",
    "        H.append(atom_feature(m[i][0]))\n",
    "    H = np.array(H)\n",
    "\n",
    "    return H\n",
    "\n",
    "pk = deepchem.dock.ConvexHullPocketFinder()\n",
    "def process_protein(pdb_file):\n",
    "    m = Chem.MolFromPDBFile(pdb_file)\n",
    "    am = GetAdjacencyMatrix(m) # n2 x n2\n",
    "    pockets = pk.find_pockets(pdb_file)\n",
    "    n2 = m.GetNumAtoms()\n",
    "    c2 = m.GetConformers()[0] # Tid\n",
    "    d2 = np.array(c2.GetPositions()) # n2 x 3\n",
    "    binding_parts = []\n",
    "    not_in_binding = [i for i in range(0, n2)]\n",
    "    constructed_graphs = []\n",
    "    for bound_box in pockets:\n",
    "        x_min = bound_box.x_range[0]\n",
    "        x_max = bound_box.x_range[1]\n",
    "        y_min = bound_box.y_range[0]\n",
    "        y_max = bound_box.y_range[1]\n",
    "        z_min = bound_box.z_range[0]\n",
    "        z_max = bound_box.z_range[1]\n",
    "        binding_parts_atoms = []\n",
    "        idxs = []\n",
    "        for idx, atom_cord in enumerate(d2):\n",
    "            if x_min < atom_cord[0] < x_max and y_min < atom_cord[1] < y_max and z_min < atom_cord[2] < z_max:\n",
    "                binding_parts_atoms.append((m.GetAtoms()[idx], atom_cord))\n",
    "                idxs.append(idx)\n",
    "                if idx in not_in_binding:\n",
    "                    not_in_binding.remove(idx)\n",
    "\n",
    "        ami = am[np.array(idxs)[:, None], np.array(idxs)] # len(idxs) x len(idxs)\n",
    "        H = get_atom_feature(binding_parts_atoms)\n",
    "        g = nx.convert_matrix.from_numpy_matrix(ami)\n",
    "        graph = dgl.from_networkx(g)\n",
    "        graph.ndata['h'] = torch.Tensor(H)\n",
    "        graph = dgl.add_self_loop(graph)\n",
    "        constructed_graphs.append(graph)\n",
    "        binding_parts.append(binding_parts_atoms)\n",
    "\n",
    "    constructed_graphs = dgl.batch(constructed_graphs)\n",
    "\n",
    "    return binding_parts, not_in_binding, constructed_graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "080feff7-d87b-4f9e-9ea7-a5fcb0bdc612",
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(42)\n",
    "\n",
    "node_featurizer = CanonicalAtomFeaturizer(atom_data_field='h') # Tid\n",
    "\n",
    "zero = np.eye(2)[1] # [0, 1]\n",
    "one = np.eye(2)[0] # [1, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4dc71c2-eb38-4be3-ab85-d680ef465b02",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"humanSeqPdb\") # 1787 x 2\n",
    "print(len(df['pdb_id'].unique())) # 1662\n",
    "\n",
    "with open(\"human_data.txt\", 'r') as fp:\n",
    "    train_raw = fp.read()\n",
    "\n",
    "constructed_graphs = \"\"\n",
    "raw_data = train_raw.split(\"\\n\") # 6729 list\n",
    "random.shuffle(raw_data)\n",
    "raw_data_train = raw_data[0: int(len(raw_data)*0.8)] # 5383\n",
    "raw_data_valid = raw_data[int(len(raw_data)*0.8): int(len(raw_data)*0.9)] # 673\n",
    "raw_data_test = raw_data[int(len(raw_data)*0.9): int(len(raw_data))] # 673\n",
    "del raw_data\n",
    "\n",
    "p_graphs = {}\n",
    "\n",
    "save_set = []\n",
    "i = 1\n",
    "# for item in raw_data_train:\n",
    "#     print(i)\n",
    "#     i += 1\n",
    "#     try:\n",
    "#         a = item.split()\n",
    "#         smile = a[0]\n",
    "#         sequence = a[1]\n",
    "#         pdb_code = df.loc[df[\"sequence\"] == sequence][\"pdb_id\"].item()[:-1] # Tid: -1 means omit the info of chain\n",
    "#         if pdb_code != \"6g5i\" and pdb_code != \"5t0j\" and pdb_code != \"5wve\":\n",
    "#             if pdb_code not in p_graphs.keys():\n",
    "#                 pdbl = PDBList()\n",
    "#                 pdbl.retrieve_pdb_file(\n",
    "#                     pdb_code, pdir='./pdbs/', overwrite=True, file_format=\"pdb\"\n",
    "#                 )\n",
    "#                 # Rename file to .pdb from .ent\n",
    "#                 os.rename(\n",
    "#                     './pdbs/' + \"pdb\" + pdb_code + \".ent\", './pdbs/' + pdb_code + \".pdb\"\n",
    "#                 )\n",
    "#                 # Assert file has been downloaded\n",
    "#                 assert any(pdb_code in s for s in os.listdir('./pdbs/'))\n",
    "#                 #print(f\"Downloaded PDB file for: {pdb_code}\")\n",
    "#                 _, _, constructed_graphs = process_protein(f\"./pdbs/{pdb_code}.pdb\")\n",
    "\n",
    "#                 p_graphs[pdb_code] = constructed_graphs\n",
    "#             else:\n",
    "#                 constructed_graphs = p_graphs[pdb_code]\n",
    "\n",
    "#             g = smiles_to_bigraph(smile, node_featurizer=node_featurizer)\n",
    "#             g = dgl.add_self_loop(g)\n",
    "#             if a[2] == \"1\":\n",
    "#                 save_set.append(((constructed_graphs, g), one))\n",
    "#             else:\n",
    "#                 save_set.append(((constructed_graphs, g), zero)) # FIXME\n",
    "#     except Exception as e:\n",
    "#         print(e)\n",
    "#         continue\n",
    "\n",
    "\n",
    "# with open(f'human_part_train.pkl', 'wb') as f:\n",
    "#     pickle.dump(save_set, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "68ee7107-0ea6-49d9-9a77-f0ed5f383ea0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading PDB structure '4iar'...\n"
     ]
    }
   ],
   "source": [
    "p_graphs = {}\n",
    "\n",
    "save_set = []\n",
    "item = 'CC[C@@]1(C[C@@H]2C3=CC(=C(C=C3CCN2C[C@H]1CC(C)C)OC)OC)O MSPLNQSAEGLPQEASNRSLNATETSEAWDPRTLQALKISLAVVLSVITLATVLSNAFVLTTILLTRKLHTPANYLIGSLATTDLLVSILVMPISIAYTITHTWNFGQILCDIWLSSDITCCTASILHLCVIALDRYWAITDALEYSKRRTAGHAATMIAIVWAISICISIPPLFWRQAKAQEEMSDCLVNTSQISYTIYSTCGAFYIPSVLLIILYGRIYRAARNRILNPPSLYGKRFTTAHLITGSAGSSLCSLNSSLHEGHSHSAGSPLFFNHVKIKLADSALERKRISAARERKATKILGIILGAFIICWLPFFVVSLVLPICRDSCWIHPALFDFFTWLGYLNSLINPIIYTVFNEEFRQAFQKIVPFRKAS 0'\n",
    "\n",
    "a = item.split()\n",
    "smile = a[0]\n",
    "sequence = a[1]\n",
    "pdb_code = '4iarA'[: 4]\n",
    "if pdb_code != \"6g5i\" and pdb_code != \"5t0j\" and pdb_code != \"5wve\":\n",
    "    if pdb_code not in p_graphs.keys():\n",
    "        pdbl = PDBList()\n",
    "        pdbl.retrieve_pdb_file(\n",
    "            pdb_code, pdir='./pdbs/', overwrite=True, file_format=\"pdb\"\n",
    "        )\n",
    "        # Rename file to .pdb from .ent\n",
    "        os.rename(\n",
    "            './pdbs/' + \"pdb\" + pdb_code + \".ent\", './pdbs/' + pdb_code + \".pdb\"\n",
    "        )\n",
    "        # Assert file has been downloaded\n",
    "        assert any(pdb_code in s for s in os.listdir('./pdbs/'))\n",
    "        #print(f\"Downloaded PDB file for: {pdb_code}\")\n",
    "        _, _, constructed_graphs = process_protein(f\"./pdbs/{pdb_code}.pdb\")\n",
    "\n",
    "        p_graphs[pdb_code] = constructed_graphs\n",
    "    else:\n",
    "        constructed_graphs = p_graphs[pdb_code]\n",
    "\n",
    "    g = smiles_to_bigraph(smile, node_featurizer=node_featurizer)\n",
    "    g = dgl.add_self_loop(g)\n",
    "    if a[2] == \"1\":\n",
    "        save_set.append(((constructed_graphs, g), one))\n",
    "    else:\n",
    "        save_set.append(((constructed_graphs, g), zero))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fc3b9079-0e7f-42c8-bc2b-6d500b18933d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'4iar': Graph(num_nodes=14086, num_edges=41200,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={})}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "949fb377-0949-416d-9f74-d94c140c5cd0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[((Graph(num_nodes=14086, num_edges=41200,\n",
       "         ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "         edata_schemes={}),\n",
       "   Graph(num_nodes=25, num_edges=79,\n",
       "         ndata_schemes={'h': Scheme(shape=(74,), dtype=torch.float32)}\n",
       "         edata_schemes={})),\n",
       "  array([0., 1.]))]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "save_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "dd35792c-3b87-4c50-b237-df7afb664c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb_ids = ['6T29_1', '6TLU_1', '6QX9_43']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "b9ceb798-2e06-449b-b76f-85f00270a011",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading PDB structure '6t29'...\n",
      "Downloading PDB structure '6tlu'...\n",
      "Downloading PDB structure '6qx9'...\n"
     ]
    }
   ],
   "source": [
    "from Bio.PDB import *\n",
    "\n",
    "for pdb_id in pdb_ids:\n",
    "    pdb_code = pdb_id[: 4].lower()\n",
    "    pdbl = PDBList()\n",
    "    pdbl.retrieve_pdb_file(\n",
    "        pdb_code, pdir='./pdbs/', overwrite=False, file_format=\"mmCif\"\n",
    "    )\n",
    "    # # Rename file to .pdb from .ent\n",
    "    # os.rename(\n",
    "    #     './pdbs/' + \"pdb\" + pdb_code + \".ent\", './pdbs/' + pdb_code + \".pdb\"\n",
    "    # )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9d84d3bd-55de-4ff2-8d1e-9bf836777046",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pdbs/6t29.cif\n",
      "./pdbs/6tlu.cif\n",
      "./pdbs/AF-P80192-F1-model_v4.cif\n",
      "./pdbs/6qx9.cif\n"
     ]
    }
   ],
   "source": [
    "for fn in os.listdir('./pdbs'):\n",
    "    if fn.endswith('.cif'):\n",
    "        print('./pdbs/' + fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "42120894-7e2d-4098-995b-2a80f6493b7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python cif2pdb.py ./pdbs/AF-P80192-F1-model_v4.cif ./pdbs/AF-P80192-F1-model_v4.pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "32ee4a33-8a03-42cb-b6f9-a456f93b10e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python cif2pdb.py ./pdbs/6t29.cif ./pdbs/6t29.pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "60c97d65-1cab-49d2-afb9-42463cc28aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python cif2pdb.py ./pdbs/6tlu.cif ./pdbs/6tlu.pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d587ae27-667c-4150-8d1d-80fba209ae86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/luozc/miniconda3/envs/FLAG/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain A2 is discontinuous at line 137419.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/FLAG/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain 5C is discontinuous at line 137420.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/FLAG/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain 5A is discontinuous at line 137453.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/FLAG/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain U is discontinuous at line 137489.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/FLAG/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain BP is discontinuous at line 137490.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/FLAG/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain 1C is discontinuous at line 137493.\n",
      "  warnings.warn(\n",
      "ERROR: Too many chains to represent in PDB format\n"
     ]
    }
   ],
   "source": [
    "!python cif2pdb.py ./pdbs/6qx9.cif ./pdbs/6qx9.pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "efacef5c-8777-4b92-b045-d83016973ea1",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'pdb_ids' 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[17], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m fn \u001b[38;5;129;01min\u001b[39;00m os\u001b[38;5;241m.\u001b[39mlistdir(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m./pdbs\u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[0;32m----> 2\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(pdb_id[: \u001b[38;5;241m4\u001b[39m]\u001b[38;5;241m.\u001b[39mlower() \u001b[38;5;129;01min\u001b[39;00m fn \u001b[38;5;28;01mfor\u001b[39;00m pdb_id \u001b[38;5;129;01min\u001b[39;00m \u001b[43mpdb_ids\u001b[49m):\n\u001b[1;32m      3\u001b[0m         \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m./pdbs/\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m fn)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'pdb_ids' is not defined"
     ]
    }
   ],
   "source": [
    "for fn in os.listdir('./pdbs'):\n",
    "    if any(pdb_id[: 4].lower() in fn for pdb_id in pdb_ids):\n",
    "        print('./pdbs/' + fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "e8515d72-b4a7-4659-b553-1c3c15a74dc7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain A2 is discontinuous at line 137419.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain 5C is discontinuous at line 137420.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain 5A is discontinuous at line 137453.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain U is discontinuous at line 137489.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain BP is discontinuous at line 137490.\n",
      "  warnings.warn(\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/Bio/PDB/StructureBuilder.py:89: PDBConstructionWarning: WARNING: Chain 1C is discontinuous at line 137493.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<Chain id=1>,\n",
       " <Chain id=6>,\n",
       " <Chain id=5O>,\n",
       " <Chain id=B4>,\n",
       " <Chain id=13>,\n",
       " <Chain id=4B>,\n",
       " <Chain id=5e>,\n",
       " <Chain id=I>,\n",
       " <Chain id=1K>,\n",
       " <Chain id=4C>,\n",
       " <Chain id=41>,\n",
       " <Chain id=R>,\n",
       " <Chain id=1f>,\n",
       " <Chain id=4e>,\n",
       " <Chain id=66>,\n",
       " <Chain id=X>,\n",
       " <Chain id=22>,\n",
       " <Chain id=5>,\n",
       " <Chain id=67>,\n",
       " <Chain id=62>,\n",
       " <Chain id=2B>,\n",
       " <Chain id=53>,\n",
       " <Chain id=A2>,\n",
       " <Chain id=B2>,\n",
       " <Chain id=2f>,\n",
       " <Chain id=5C>,\n",
       " <Chain id=5X>,\n",
       " <Chain id=11>,\n",
       " <Chain id=42>,\n",
       " <Chain id=1b>,\n",
       " <Chain id=B5>,\n",
       " <Chain id=1A>,\n",
       " <Chain id=S>,\n",
       " <Chain id=5f>,\n",
       " <Chain id=5J>,\n",
       " <Chain id=51>,\n",
       " <Chain id=4D>,\n",
       " <Chain id=63>,\n",
       " <Chain id=2b>,\n",
       " <Chain id=2>,\n",
       " <Chain id=4f>,\n",
       " <Chain id=B3>,\n",
       " <Chain id=1g>,\n",
       " <Chain id=23>,\n",
       " <Chain id=5b>,\n",
       " <Chain id=68>,\n",
       " <Chain id=43>,\n",
       " <Chain id=1e>,\n",
       " <Chain id=5A>,\n",
       " <Chain id=A3>,\n",
       " <Chain id=U>,\n",
       " <Chain id=2g>,\n",
       " <Chain id=5D>,\n",
       " <Chain id=52>,\n",
       " <Chain id=12>,\n",
       " <Chain id=64>,\n",
       " <Chain id=2e>,\n",
       " <Chain id=BP>,\n",
       " <Chain id=1C>,\n",
       " <Chain id=5g>,\n",
       " <Chain id=K>,\n",
       " <Chain id=4b>,\n",
       " <Chain id=4A>,\n",
       " <Chain id=21>,\n",
       " <Chain id=4g>,\n",
       " <Chain id=4>,\n",
       " <Chain id=2A>,\n",
       " <Chain id=A1>,\n",
       " <Chain id=65>,\n",
       " <Chain id=5B>,\n",
       " <Chain id=B1>]"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from Bio.PDB.MMCIFParser import MMCIFParser\n",
    "\n",
    "ciffile = './pdbs/6qx9.cif'\n",
    "strucid = ciffile[:4] if len(ciffile)>4 else \"1xxx\"\n",
    "parser = MMCIFParser()\n",
    "structure = parser.get_structure(strucid, ciffile)\n",
    "list(structure.get_chains())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "0b6c2997-1623-4106-b6f4-492cc4772580",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "71"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(list(structure.get_chains()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "cca3a937-4974-48e6-b182-afb060e920de",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8,\n",
       " ['7 PDB id: 3MY0_1 ',\n",
       "  \"27 PDB id: 7MFE_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"34 PDB id: 6W4P_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"46 PDB id: 1UA2_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"67 PDB id: 5OOI_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"87 PDB id: 4BKA_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"88 PDB id: 4M4R_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"93 PDB id: 7MN5_3 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\"])"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./davis_1_site_msg.pkl', 'rb') as f:\n",
    "    msg = pickle.load(f)\n",
    "len(msg), msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "dde40f84-f712-469c-92d4-1134c80126df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9,\n",
       " [\"8 PDB id: 5T89_2 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"9 PDB id: 7QDP_2 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"34 PDB id: 4Z32_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"38 PDB id: 5NCL_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"58 PDB id: 3VUT_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  '76 PDB id: 5ZCS_1 ',\n",
       "  \"79 PDB id: 5I0I_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"83 PDB id: 6TM5_17 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\",\n",
       "  \"85 PDB id: 5CI6_1 Python argument types in\\n    rdkit.Chem.rdmolops.GetAdjacencyMatrix(NoneType)\\ndid not match C++ signature:\\n    GetAdjacencyMatrix(RDKit::ROMol {lvalue} mol, bool useBO=False, int emptyVal=0, bool force=False, char const* prefix='')\"])"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./davis_2_site_msg.pkl', 'rb') as f:\n",
    "    msg = pickle.load(f)\n",
    "len(msg), msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c019885b-161f-43b4-8ee4-59d9c28a4476",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df_davis.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "76903c96-d961-48dd-9c28-21353364e4e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-02-06 09:04:34.143489: 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",
      "2023-02-06 09:04:34.867119: 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",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/tensorflow/python/framework/dtypes.py:246: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  np.bool8: (False, True),\n",
      "2023-02-06 09:04:38.752193: 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; LD_LIBRARY_PATH: /home/luozc/CCMpred/lib:/home/luozc/miniconda3/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:\n",
      "2023-02-06 09:04:38.752476: 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; LD_LIBRARY_PATH: /home/luozc/CCMpred/lib:/home/luozc/miniconda3/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:/usr/lib/x86_64-linux-gnu/:/home/luozc/miniconda3/lib:/usr/local/cuda-10.0/lib64:/home/luozc/suitesparse/lib:/home/luozc/MPFR/lib:/home/luozc/GMP/lib:/home/luozc/suitesparse/SuiteSparse/lib:/home/luozc/lapack:/home/luozc/BLASandLAPACK/lib:/home/luozc/gflags/lib:/home/luozc/google-glog/lib:/home/luozc/openssl/lib:\n",
      "2023-02-06 09:04:38.752489: 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/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/torch/utils/tensorboard/__init__.py:4: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if not hasattr(tensorboard, \"__version__\") or LooseVersion(\n",
      "/home/luozc/miniconda3/envs/pl102/lib/python3.10/site-packages/torch/utils/tensorboard/__init__.py:6: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  ) < LooseVersion(\"1.15\"):\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "from collections import OrderedDict\n",
    "import re\n",
    "import csv\n",
    "import pypdb\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "\n",
    "import pandas as pd\n",
    "import dgl\n",
    "import numpy as np\n",
    "import torch\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem.rdmolops import GetAdjacencyMatrix\n",
    "import networkx as nx\n",
    "from Bio.PDB import *\n",
    "import deepchem\n",
    "import pickle\n",
    "import time\n",
    "\n",
    "def one_of_k_encoding(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        raise Exception(\"input {0} not in allowable set{1}:\".format(x, allowable_set))\n",
    "    # print list((map(lambda s: x == s, allowable_set)))\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def one_of_k_encoding_unk(x, allowable_set):\n",
    "    \"\"\"Maps inputs not in the allowable set to the last element.\"\"\"\n",
    "    if x not in allowable_set:\n",
    "        x = allowable_set[-1]\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def atom_feature(atom):\n",
    "    return np.array(one_of_k_encoding_unk(atom.GetSymbol(),\n",
    "                                          ['C', 'N', 'O', 'S', 'F', 'P', 'Cl', 'Br', 'B', 'H']) +\n",
    "                    one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5, 6, 7, 8]) +\n",
    "                    one_of_k_encoding_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4]) +\n",
    "                    one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5]) +\n",
    "                    [atom.GetIsAromatic()])  # (10, 9, 5, 6, 1) --> total 31\n",
    "\n",
    "\n",
    "def get_atom_feature(m):\n",
    "    H = []\n",
    "    for i in range(len(m)):\n",
    "        H.append(atom_feature(m[i][0]))\n",
    "    H = np.array(H)\n",
    "\n",
    "    return H\n",
    "\n",
    "pk = deepchem.dock.ConvexHullPocketFinder()\n",
    "def process_protein(pdb_file):\n",
    "    m = Chem.MolFromPDBFile(pdb_file)\n",
    "    if m is None:\n",
    "        return None, None, None\n",
    "    am = GetAdjacencyMatrix(m) # n2 x n2\n",
    "    pockets = pk.find_pockets(pdb_file)\n",
    "    n2 = m.GetNumAtoms()\n",
    "    c2 = m.GetConformers()[0] # Tid\n",
    "    d2 = np.array(c2.GetPositions()) # n2 x 3\n",
    "    binding_parts = []\n",
    "    not_in_binding = [i for i in range(0, n2)]\n",
    "    constructed_graphs = []\n",
    "    for bound_box in pockets:\n",
    "        x_min = bound_box.x_range[0]\n",
    "        x_max = bound_box.x_range[1]\n",
    "        y_min = bound_box.y_range[0]\n",
    "        y_max = bound_box.y_range[1]\n",
    "        z_min = bound_box.z_range[0]\n",
    "        z_max = bound_box.z_range[1]\n",
    "        binding_parts_atoms = []\n",
    "        idxs = []\n",
    "        for idx, atom_cord in enumerate(d2):\n",
    "            if x_min < atom_cord[0] < x_max and y_min < atom_cord[1] < y_max and z_min < atom_cord[2] < z_max:\n",
    "                binding_parts_atoms.append((m.GetAtoms()[idx], atom_cord))\n",
    "                idxs.append(idx)\n",
    "                if idx in not_in_binding:\n",
    "                    not_in_binding.remove(idx)\n",
    "\n",
    "        ami = am[np.array(idxs)[:, None], np.array(idxs)] # len(idxs) x len(idxs)\n",
    "        H = get_atom_feature(binding_parts_atoms)\n",
    "        g = nx.convert_matrix.from_numpy_matrix(ami)\n",
    "        graph = dgl.from_networkx(g)\n",
    "        graph.ndata['h'] = torch.Tensor(H)\n",
    "        graph = dgl.add_self_loop(graph)\n",
    "        constructed_graphs.append(graph)\n",
    "        binding_parts.append(binding_parts_atoms)\n",
    "\n",
    "    constructed_graphs = dgl.batch(constructed_graphs)\n",
    "\n",
    "    return binding_parts, not_in_binding, constructed_graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8635a9fa-b6fe-4ffc-91a0-b98b08b0f2d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['6GL7_3', '6ZXF_35', '6G51_23', '6G5I_36', '4NEU_1'], dtype=object)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['target_pdb_id'].unique()[239: 244]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0f3a50f5-0f1d-4439-836b-44cbedd8d1ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3,\n",
       " ['7 PDB id: 3MY0_1 ',\n",
       "  '176 PDB id: 5ZCS_1 ',\n",
       "  '236 PDB id: 6QX9_43 Bad input file ./pdbs/6qx9.pdb'])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./davis_site_msg.pkl', 'rb') as f:\n",
    "    msg = pickle.load(f)\n",
    "len(msg), msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "adf44e28-8f41-4b2b-a74f-4170641baaec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, ['81 PDB id: 5ZCS_1 '])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./kiba_site_msg.pkl', 'rb') as f:\n",
    "    msg = pickle.load(f)\n",
    "len(msg), msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "48a1a11c-12c1-4874-925c-4d564c745eb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['QLTPTNSLKRGGAHHRRCEVALLGCGAVLAATGLGFDLLEAGKCQLLPLEEPEPPAREEKKRREGLFQRSSRPRRSTSPPSRKLFKKEEPMLLLGDPSASLTLLSLSSISECNSTRSLLRSDSDEIVVYEMPVSPVEAPPLSPCTHNPLVNVRVERFKRDPNQSLTPTHVTLTTPSQPSSHRRTPSDGALKPETLLASRSPSSNGLSPSPGAGMLKTPSPSRDPGEFPRLPDPNVVFPPTPRRWNTQQDSTLERPKTLEFLPRPRPSANRQRLDPWWFVSPSHARSTSPANSSSTETPSNLDSCFASSSSTVEERPGLPALLPFQAGPLPPTERTLLDLDAEGQSQDSTVPLCRAELNTHRPAPYEIQQEFWS'],\n",
       "      dtype=object)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = 'davis'\n",
    "df_davis[df_davis['target_pdb_id'] == 'TOFOLD']['target_sequence'].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "45056631-4e83-4002-b35f-bff594d803cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "prot_seq_AF_dict = {'QLTPTNSLKRGGAHHRRCEVALLGCGAVLAATGLGFDLLEAGKCQLLPLEEPEPPAREEKKRREGLFQRSSRPRRSTSPPSRKLFKKEEPMLLLGDPSASLTLLSLSSISECNSTRSLLRSDSDEIVVYEMPVSPVEAPPLSPCTHNPLVNVRVERFKRDPNQSLTPTHVTLTTPSQPSSHRRTPSDGALKPETLLASRSPSSNGLSPSPGAGMLKTPSPSRDPGEFPRLPDPNVVFPPTPRRWNTQQDSTLERPKTLEFLPRPRPSANRQRLDPWWFVSPSHARSTSPANSSSTETPSNLDSCFASSSSTVEERPGLPALLPFQAGPLPPTERTLLDLDAEGQSQDSTVPLCRAELNTHRPAPYEIQQEFWS':'AF-P80192-F1-model_v4'}\n",
    "\n",
    "with open(dataset + '_seq_AF.csv', 'w') as f:\n",
    "    f.write('target_sequence,target_AF_id\\n')\n",
    "    for seq, AF_id in prot_seq_AF_dict.items():\n",
    "        ls = []\n",
    "        ls += [ seq ]\n",
    "        ls += [ AF_id ]\n",
    "        f.write(','.join(map(str, ls)) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "db521549-c833-491e-abb6-8fb2442ca124",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "292 <class 'dict'>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.73s/it]\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(dataset + '_seq_AF.csv')\n",
    "\n",
    "with open(f'{dataset}_site.pkl', 'rb') as f:\n",
    "    p_graphs = pickle.load(f)\n",
    "print(len(p_graphs), type(p_graphs))\n",
    "\n",
    "for AF_id in tqdm(df['target_AF_id'].unique()):\n",
    "    try:\n",
    "        AF_code = AF_id\n",
    "        if AF_code not in p_graphs.keys():\n",
    "            # Assert file has been downloaded\n",
    "            assert any(AF_code in s for s in os.listdir('./pdbs/'))\n",
    "            #print(f\"Downloaded PDB file for: {AF_code}\")\n",
    "            _, _, constructed_graphs = process_protein(f\"./pdbs/{AF_code}.pdb\")\n",
    "\n",
    "            p_graphs[AF_code] = constructed_graphs\n",
    "    except Exception as e:\n",
    "        print(f'{time.strftime(\"%Y-%m-%d|%H:%M:%S\", time.localtime())}: AF_id={AF_id}')\n",
    "        continue\n",
    "\n",
    "with open(f'{dataset}_site.pkl', 'wb') as f:\n",
    "    pickle.dump(p_graphs, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c54b69e8-466b-401b-9829-d121588b368b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{dataset}_site.pkl', 'rb') as f:\n",
    "    p_graphs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d10a41bc-521f-47d3-9919-dc5763e9854e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cnt = 0\n",
    "for val in p_graphs.values():\n",
    "    if val is None:\n",
    "        cnt += 1\n",
    "cnt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8458beb2-cc9f-41ca-a503-863e535d443e",
   "metadata": {},
   "source": [
    "# Filter tasks without Hybrid Graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "71efabee-8831-42b9-99a7-0d97dfa98ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prot_key(row):\n",
    "    pdb_id = seq2pdb[seq2pdb['target_sequence'] == row['target_sequence']]['target_pdb_id'].to_list()[0]\n",
    "    if pdb_id == 'TOFOLD':\n",
    "        return seq2AF[seq2AF['target_sequence'] == row['target_sequence']]['target_AF_id'].to_list()[0]\n",
    "    else:\n",
    "        return pdb_id[: 4].lower()\n",
    "\n",
    "def add_prot_key_col(df):\n",
    "    new_df = df.copy()\n",
    "    new_df['prot_key'] = df.apply(get_prot_key, axis=1)\n",
    "    return new_df\n",
    "\n",
    "def filter_by_prot_key(df):\n",
    "    new_df = df.copy()\n",
    "    return new_df[new_df.apply(lambda row: row['prot_key'] in prot_key2prot_graph.keys(), axis=1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "55127033-89e8-41cb-a58f-ceb85e653746",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['compound_iso_smiles', 'target_sequence', 'target_pdb_id', 'affinity',\n",
       "       'protein_id', 'drug_id'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "99eb1365-2226-4c8f-a0ac-6de7100e1080",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|███████████████████████                                                                                            | 1/5 [00:09<00:39,  9.79s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 0 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|██████████████████████████████████████████████                                                                     | 2/5 [00:19<00:29,  9.85s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 1 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|█████████████████████████████████████████████████████████████████████                                              | 3/5 [00:29<00:19,  9.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 2 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████████████████████████████████████████████████████████████████████████████████████████                       | 4/5 [00:39<00:09,  9.91s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 3 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:49<00:00,  9.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 4 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      " 20%|███████████████████████                                                                                            | 1/5 [00:35<02:23, 35.75s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 0 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|██████████████████████████████████████████████                                                                     | 2/5 [01:12<01:48, 36.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 1 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|█████████████████████████████████████████████████████████████████████                                              | 3/5 [01:47<01:11, 35.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 2 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████████████████████████████████████████████████████████████████████████████████████████                       | 4/5 [02:23<00:35, 36.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 3 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [02:59<00:00, 36.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 4 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "datasets = ['davis', 'kiba']\n",
    "folds = '01234'\n",
    "\n",
    "for dataset in datasets:\n",
    "    # test df\n",
    "    test_df = pd.read_csv(f'data/{dataset}/{dataset}_test.csv')\n",
    "    \n",
    "    # filter no Hybrid graphs task\n",
    "    if dataset == 'davis':\n",
    "        seq2AF = pd.read_csv(f'{dataset}_seq_AF.csv')\n",
    "    seq2pdb = pd.read_csv(f'{dataset}_seq_pdb.csv')\n",
    "    \n",
    "    with open(f'{dataset}_site.pkl', 'rb') as fp:\n",
    "        prot_key2prot_graph = pickle.load(fp)\n",
    "    prot_key2prot_graph = {key: val for key, val in prot_key2prot_graph.items() if val is not None}\n",
    "    \n",
    "    filter_by_prot_key(add_prot_key_col(test_df)).to_csv(f'data/{dataset}/{dataset}_test_clr.csv', index=False)\n",
    "    \n",
    "    for fold in tqdm(folds):\n",
    "        # train & val df\n",
    "        train_df = pd.read_csv(f'data/{dataset}/{dataset}_train_fold_{fold}.csv')\n",
    "        valid_df = pd.read_csv(f'data/{dataset}/{dataset}_valid_fold_{fold}.csv')\n",
    "\n",
    "        filter_by_prot_key(add_prot_key_col(train_df)).to_csv(f'data/{dataset}/{dataset}_train_fold_{fold}_clr.csv', index=False)\n",
    "        filter_by_prot_key(add_prot_key_col(valid_df)).to_csv(f'data/{dataset}/{dataset}_valid_fold_{fold}_clr.csv', index=False)\n",
    "        \n",
    "        print(dataset, f'fold {fold} clear!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "a5f35659-2384-44a3-a323-ed57878750fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis cold clear!\n",
      "kiba cold clear!\n"
     ]
    }
   ],
   "source": [
    "datasets = ['davis', 'kiba']\n",
    "\n",
    "for dataset in datasets:    \n",
    "    # filter no Hybrid graphs task\n",
    "    if dataset == 'davis':\n",
    "        seq2AF = pd.read_csv(f'{dataset}_seq_AF.csv')\n",
    "    seq2pdb = pd.read_csv(f'{dataset}_seq_pdb.csv')\n",
    "    \n",
    "    with open(f'{dataset}_site.pkl', 'rb') as fp:\n",
    "        prot_key2prot_graph = pickle.load(fp)\n",
    "    prot_key2prot_graph = {key: val for key, val in prot_key2prot_graph.items() if val is not None}\n",
    "    \n",
    "    # cold df\n",
    "    cold_df = pd.read_csv(f'data/{dataset}_cold.csv')\n",
    "\n",
    "    filter_by_prot_key(add_prot_key_col(cold_df)).to_csv(f'data/{dataset}_cold_clr.csv', index=False)\n",
    "    \n",
    "    print(dataset, f'cold clear!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "69b1925c-6812-49bb-88bd-482ddea0581d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def atom_features(atom):\n",
    "    # 44 + 11 + 11 + 11 + 1 + 3 + 1\n",
    "    return np.array(one_of_k_encoding_unk(atom.GetSymbol(),\n",
    "                                          ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca', 'Fe', 'As',\n",
    "                                           'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag', 'Pd', 'Co', 'Se',\n",
    "                                           'Ti', 'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd', 'In', 'Mn', 'Zr', 'Cr',\n",
    "                                           'Pt', 'Hg', 'Pb', 'X']) +\n",
    "                    one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +\n",
    "                    one_of_k_encoding_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +\n",
    "                    one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +\n",
    "                    [atom.GetIsAromatic()]\n",
    "                    +\n",
    "                    one_of_k_encoding_unk(atom.GetFormalCharge() , [-1,0,1]) +\n",
    "                    [atom.IsInRing()]\n",
    ")\n",
    "\n",
    "\n",
    "# one ont encoding\n",
    "def one_of_k_encoding(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        # print(x)\n",
    "        raise Exception('input {0} not in allowable set{1}:'.format(x, allowable_set))\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def one_of_k_encoding_unk(x, allowable_set):\n",
    "    '''Maps inputs not in the allowable set to the last element.'''\n",
    "    if x not in allowable_set:\n",
    "        x = allowable_set[-1]\n",
    "    return list(map(lambda s: x == s, allowable_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "61e69703-cce6-49ab-a213-fd350363bd53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drug smiles -> Graph, call func-smile_to_graph, exist for sure\n",
    "def smile_to_graph(smile):\n",
    "    mol = Chem.MolFromSmiles(smile)\n",
    "    if mol is None:\n",
    "        return None\n",
    "    c_size = mol.GetNumAtoms()\n",
    "\n",
    "    features = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        feature = atom_features(atom)\n",
    "        features.append(feature / sum(feature))\n",
    "\n",
    "    edges = []\n",
    "    for bond in mol.GetBonds():\n",
    "        edges.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])\n",
    "    g = nx.Graph(edges).to_directed()\n",
    "    edge_index = []\n",
    "    mol_adj = np.zeros((c_size, c_size))\n",
    "    for e1, e2 in g.edges:\n",
    "        mol_adj[e1, e2] = 1\n",
    "        # edge_index.append([e1, e2])\n",
    "    mol_adj += np.matrix(np.eye(mol_adj.shape[0]))\n",
    "    index_row, index_col = np.where(mol_adj >= 0.5)\n",
    "    for i, j in zip(index_row, index_col):\n",
    "        edge_index.append([i, j])\n",
    "    return c_size, features, edge_index\n",
    "\n",
    "# Prot seq -> LLM, use pid2llm dict, exist for sure\n",
    "pid2llm = np.load(f'{dataset}.npz',allow_pickle=True)['dict'][()]\n",
    "def pid_to_llm(pid):\n",
    "    return torch.from_numpy(pid2llm[pid]).squeeze()\n",
    "\n",
    "# Prot seq -> Hybrid graphs, seq2[AF | pdb_id[: 4]], then use site.pkl, may not exist(no [pdb_id] or None)\n",
    "with open(f'{dataset}_site.pkl', 'rb') as fp:\n",
    "    prot_key2prot_graph = pickle.load(fp)\n",
    "def prot_key_to_graph(prot_key):\n",
    "    return prot_key2prot_graph[prot_key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "73829636-c364-474a-8962-c7c2e9a815ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{dataset}_site.pkl', 'rb') as fp:\n",
    "    prot_key2prot_graph = pickle.load(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "a7b625cc-15bb-4612-8a78-ed64968ee972",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{dataset}_site.pkl', 'rb') as fp:\n",
    "    prot_key2prot_graph = pickle.load(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "0f3e2309-5f42-4f1c-9f93-50d882d1b0b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'davis'\n",
    "fold = '0'\n",
    "\n",
    "train_df = pd.read_csv(f'data/{dataset}/{dataset}_train_fold_{fold}_clr.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "d9643f43-52eb-45c0-989f-1bb4a79d01cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df_sub = train_df[: 10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "ad7fe864-cc13-4675-b2c1-db855cc6a001",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([19, 82]) torch.Size([61, 2]) torch.Size([849, 2560])\n",
      "Graph(num_nodes=12701, num_edges=36429,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([35, 82]) torch.Size([119, 2]) torch.Size([774, 2560])\n",
      "Graph(num_nodes=16369, num_edges=45211,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([34, 82]) torch.Size([108, 2]) torch.Size([495, 2560])\n",
      "Graph(num_nodes=15888, num_edges=42868,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([37, 82]) torch.Size([117, 2]) torch.Size([526, 2560])\n",
      "Graph(num_nodes=19148, num_edges=54038,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([36, 82]) torch.Size([116, 2]) torch.Size([589, 2560])\n",
      "Graph(num_nodes=41359, num_edges=119859,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([25, 82]) torch.Size([79, 2]) torch.Size([1024, 2560])\n",
      "Graph(num_nodes=19823, num_edges=55623,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([25, 82]) torch.Size([79, 2]) torch.Size([747, 2560])\n",
      "Graph(num_nodes=44729, num_edges=131149,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([46, 82]) torch.Size([148, 2]) torch.Size([841, 2560])\n",
      "Graph(num_nodes=13191, num_edges=35285,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([38, 82]) torch.Size([122, 2]) torch.Size([768, 2560])\n",
      "Graph(num_nodes=21524, num_edges=59616,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n",
      "torch.Size([33, 82]) torch.Size([105, 2]) torch.Size([640, 2560])\n",
      "Graph(num_nodes=24118, num_edges=66726,\n",
      "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
      "      edata_schemes={})\n"
     ]
    }
   ],
   "source": [
    "for idx, row in train_df_sub.iterrows():\n",
    "    cis, pid, key = row.compound_iso_smiles, row.protein_id, row.prot_key\n",
    "    mol_g = smile_to_graph(cis)\n",
    "    prot_llm = pid_to_llm(pid)\n",
    "    prot_g = prot_key_to_graph(key)\n",
    "    \n",
    "    print(torch.Tensor(mol_g[1]).shape, torch.Tensor(mol_g[2]).shape, prot_llm.shape)\n",
    "    print(prot_g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "fada919c-a103-4402-8201-fe6ad8b5a4ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Graph(num_nodes=24118, num_edges=66726,\n",
       "      ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "      edata_schemes={})"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prot_g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "024f1c89-a157-4acf-bf81-ae291d4a6b41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([24118, 31])"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(prot_g.ndata['h']).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "d0ac5a5f-2312-4b61-80a3-631dfdc18bca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_graph': <dgl.heterograph_index.HeteroGraphIndex at 0x7f97f4c8a640>,\n",
       " '_canonical_etypes': [('_N', '_E', '_N')],\n",
       " '_batch_num_nodes': {'_N': tensor([ 174,  141,  860,  435,  457,  518,  324,  413,  283,  236,  144,  195,\n",
       "           248,  461,  484,  596,  204,  352,  497,  577,  173,  161,  648,  263,\n",
       "           631,  768,   96,   55,  146,  257, 1482,  516,  287,  242,  214,  421,\n",
       "           335,  443,  176,  315,  252,   87,  106,  259,  372,  170,  493,  342,\n",
       "           653,  573,  478,  255,  300,  427,  392,  302,   77,  223,  232,  500,\n",
       "           106,  881,  313,  203,  257,  637])},\n",
       " '_batch_num_edges': {('_N',\n",
       "   '_E',\n",
       "   '_N'): tensor([ 484,  399, 2412, 1211, 1297, 1454,  922, 1115,  797,  598,  392,  531,\n",
       "           678, 1255, 1330, 1640,  596,  998, 1369, 1605,  471,  439, 1814,  741,\n",
       "          1787, 2150,  266,  147,  390,  725, 4202, 1454,  799,  676,  600, 1153,\n",
       "           921, 1211,  486,  869,  688,  241,  288,  717, 1020,  482, 1337,  940,\n",
       "          1799, 1515, 1276,  707,  836, 1181, 1054,  824,  215,  619,  584, 1344,\n",
       "           300, 2501,  863,  561,  729, 1721])},\n",
       " '_ntypes': ['_N'],\n",
       " '_is_unibipartite': False,\n",
       " '_srctypes_invmap': {'_N': 0},\n",
       " '_dsttypes_invmap': {'_N': 0},\n",
       " '_etypes': ['_E'],\n",
       " '_etype2canonical': {'_E': ('_N', '_E', '_N')},\n",
       " '_etypes_invmap': {('_N', '_E', '_N'): 0},\n",
       " '_node_frames': [{'h': tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 0., 0.],\n",
       "          [0., 1., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 1.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 0., 0.]])}],\n",
       " '_edge_frames': [{}]}"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vars(prot_g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "ac922b3a-0e2e-4965-91e5-d3fd473d80ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['_graph', '_canonical_etypes', '_batch_num_nodes', '_batch_num_edges', '_ntypes', '_is_unibipartite', '_srctypes_invmap', '_dsttypes_invmap', '_etypes', '_etype2canonical', '_etypes_invmap', '_node_frames', '_edge_frames'])"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vars(prot_g).keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "e61beee1-fd38-4f50-96e2-359e925cfa52",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'h': tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 1.,  ..., 0., 0., 0.],\n",
       "        [0., 1., 0.,  ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 1.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 1.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 1.,  ..., 0., 0., 0.]])}"
      ]
     },
     "execution_count": 139,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prot_g.ndata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "6e567837-d09c-4bea-96f6-d9512fae176d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dgl.nn.pytorch.conv import GATConv, GraphConv, TAGConv, GINConv, APPNPConv\n",
    "from dgl.nn import TWIRLSConv\n",
    "from dgl.nn.pytorch.glob import MaxPooling, GlobalAttentionPooling\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "153d2388-0aac-4839-a94e-d4f8b455893c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prot_conv = nn.ModuleList()\n",
    "for _ in range(5):\n",
    "    prot_conv.append(TAGConv(31, 31, 2))\n",
    "\n",
    "pooling_prot = nn.Linear(31, 1)\n",
    "pool_prot = GlobalAttentionPooling(pooling_prot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "8aa8c4f8-065b-41f9-84a1-440b48d2a323",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ModuleList(\n",
       "  (0): TAGConv(\n",
       "    (lin): Linear(in_features=93, out_features=31, bias=True)\n",
       "  )\n",
       "  (1): TAGConv(\n",
       "    (lin): Linear(in_features=93, out_features=31, bias=True)\n",
       "  )\n",
       "  (2): TAGConv(\n",
       "    (lin): Linear(in_features=93, out_features=31, bias=True)\n",
       "  )\n",
       "  (3): TAGConv(\n",
       "    (lin): Linear(in_features=93, out_features=31, bias=True)\n",
       "  )\n",
       "  (4): TAGConv(\n",
       "    (lin): Linear(in_features=93, out_features=31, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 143,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prot_conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "4a0a7f1f-76ee-48cc-af43-82afce1e0006",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GlobalAttentionPooling(\n",
       "  (gate_nn): Linear(in_features=31, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pool_prot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "576557b1-22eb-4b8f-9f9c-12dfc90d8180",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([24118, 31])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([24118, 31])"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# single graph\n",
    "\n",
    "feat_prot = prot_g.ndata['h']\n",
    "print(feat_prot.shape)\n",
    "for module in prot_conv:\n",
    "    feat_prot = F.relu(module(prot_g, feat_prot))\n",
    "feat_prot.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "aae72830-3782-45f4-aaa9-438eee63d6ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([66, 31])"
      ]
     },
     "execution_count": 150,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# single graph\n",
    "\n",
    "prot_repr = pool_prot(prot_g, feat_prot)\n",
    "prot_repr.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "c7edd162-cd0f-4ed5-b23e-112c0c47cdf4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 66, 31])"
      ]
     },
     "execution_count": 152,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prot_repr = prot_repr.view(1, -1, 31)\n",
    "prot_repr.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "7fb5c71c-e7b2-4a0b-8193-e85e6c9bb9be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch graph\n",
    "\n",
    "prot_gs = []\n",
    "for idx, row in train_df_sub.iterrows():\n",
    "    key = row.prot_key\n",
    "    prot_g = prot_key_to_graph(key)\n",
    "    prot_gs.append(prot_g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "d7a139ac-60f2-4ef7-b599-1530012f2a20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Graph(num_nodes=12701, num_edges=36429,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=16369, num_edges=45211,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=15888, num_edges=42868,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=19148, num_edges=54038,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=41359, num_edges=119859,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=19823, num_edges=55623,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=44729, num_edges=131149,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=13191, num_edges=35285,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=21524, num_edges=59616,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={}),\n",
       " Graph(num_nodes=24118, num_edges=66726,\n",
       "       ndata_schemes={'h': Scheme(shape=(31,), dtype=torch.float32)}\n",
       "       edata_schemes={})]"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prot_gs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "bfdad4b3-c578-4a83-97dd-81c06d240882",
   "metadata": {},
   "outputs": [],
   "source": [
    "prot_convs = prot_conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "id": "7d73253f-ea46-4493-8bbe-938e4a32acbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "relu = nn.ReLU()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "id": "ae221039-883d-4815-8956-443da1fb9ffc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n",
      "torch.Size([1, 140, 31])\n"
     ]
    }
   ],
   "source": [
    "for prot_g in prot_gs:\n",
    "    feat_prot = prot_g.ndata['h']\n",
    "\n",
    "    for prot_conv in prot_convs:\n",
    "        feat_prot = relu(prot_conv(prot_g, feat_prot))\n",
    "\n",
    "    prot_repr = pool_prot(prot_g, feat_prot).view(1, -1, 31)\n",
    "    prot_repr = F.pad(\n",
    "        input=prot_repr, \n",
    "        pad=(0, 0, 0, 140 - prot_repr.size()[1]), \n",
    "        mode='constant', value=0)\n",
    "    print(prot_repr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "id": "a0fa231c-2abd-4889-97a1-068ae3ae262c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prot_g_bilstm = nn.LSTM(31, 31, num_layers=2, batch_first=True, bidirectional=True, dropout=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "bf780e34-d30c-431b-b0d8-9a02a448d95a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([140, 10, 62])"
      ]
     },
     "execution_count": 185,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prot_reprs = []\n",
    "for prot_g in prot_gs:\n",
    "    feat_prot = prot_g.ndata['h']\n",
    "\n",
    "    for prot_conv in prot_convs:\n",
    "        feat_prot = relu(prot_conv(prot_g, feat_prot))\n",
    "\n",
    "    prot_repr = pool_prot(prot_g, feat_prot).view(1, -1, 31)\n",
    "    prot_repr = F.pad(\n",
    "        input=prot_repr, \n",
    "        pad=(0, 0, 0, 140 - prot_repr.size()[1]), \n",
    "        mode='constant', value=0)\n",
    "    prot_repr, _ = prot_g_bilstm(prot_repr)\n",
    "    prot_reprs.append(prot_repr)\n",
    "torch.concat(prot_reprs).permute(1, 0, 2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3b44db3-b7f7-45a6-a7d0-0541ca225628",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol_gs = []\n",
    "for idx, row in train_df_sub.iterrows():\n",
    "    cis = row.compound_iso_smiles\n",
    "    mol_g = smile_to_graph(cis)\n",
    "    mol_gs.append(mol_g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "43dcc9e0-c6ae-487d-afde-3f05fe92f71b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([849, 2560])\n",
      "torch.Size([774, 2560])\n",
      "torch.Size([495, 2560])\n",
      "torch.Size([526, 2560])\n",
      "torch.Size([589, 2560])\n",
      "torch.Size([1024, 2560])\n",
      "torch.Size([747, 2560])\n",
      "torch.Size([841, 2560])\n",
      "torch.Size([768, 2560])\n",
      "torch.Size([640, 2560])\n"
     ]
    }
   ],
   "source": [
    "for idx, row in train_df_sub.iterrows():\n",
    "    cis, pid, key = row.compound_iso_smiles, row.protein_id, row.prot_key\n",
    "    prot_llm = pid_to_llm(pid)\n",
    "    print(prot_llm.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "id": "cd0cd0c2-9d5d-4e80-8037-5bf73388b8e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.TGCA.tgt_guided_cross_attention_model import TargetGuidedCrossAttention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "6e9f3ba7-cfe4-4fdc-a4ff-3757e12003ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "tgca = TargetGuidedCrossAttention(embed_dim=128, num_heads=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "87db4f37-0fec-4233-aefc-5bf016f07709",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TargetGuidedCrossAttention(\n",
       "  (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 188,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tgca"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "id": "0571c3dc-d145-4b1a-95b3-f990ce257c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "llms = torch.ones(211, 1024, 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "771cb2da-b624-4488-93db-9c410f45cb88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([211, 1024, 512])"
      ]
     },
     "execution_count": 193,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llms.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "5a591510-8067-42ca-9077-333cae3a28fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([211, 1024, 512])"
      ]
     },
     "execution_count": 199,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llms.flatten(2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "ab9d4a0e-e769-4ad5-a22c-cb80dae5a964",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([211, 512, 1024])"
      ]
     },
     "execution_count": 202,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llms.transpose(-1, -2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "id": "283b4070-8b75-44a1-8bb0-5cb9c58eac8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024\n",
      "1024\n"
     ]
    }
   ],
   "source": [
    "datasets = ['kiba', 'davis']\n",
    "\n",
    "for dataset in datasets:\n",
    "    pid2llm = np.load(f'{dataset}.npz',allow_pickle=True)['dict'][()]\n",
    "    print(max([llm.shape[1] for llm in pid2llm.values()]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f630a39-81dc-432f-b1bc-daad253daf6d",
   "metadata": {},
   "source": [
    "# CM csv clr..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0f772894-e0f3-4ed5-9aa1-585cf9dc57a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/data5/luozc/projects/DTA/GINCM-DTA/data/'\n",
    "def filter_by_tgt_ori_key(df, dataset):\n",
    "    new_df = df.copy()\n",
    "    return new_df[new_df.apply(lambda row: valid_target(row.target_original_key, dataset, data_dir), axis=1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6dbe9f9c-0a5b-4206-954a-c362f3c5c214",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'kiba'\n",
    "df = pd.read_csv(f'data/{dataset}/{dataset}_test.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146a330d-33ac-4d3d-8e1b-85821d9f09c6",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>compound_iso_smiles</th>\n",
       "      <th>target_sequence</th>\n",
       "      <th>target_pdb_id</th>\n",
       "      <th>target_original_key</th>\n",
       "      <th>affinity</th>\n",
       "      <th>protein_id</th>\n",
       "      <th>drug_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cc1ccccc1NC(=O)Nc1ccc(-c2coc3ncnc(N)c23)cc1</td>\n",
       "      <td>MPALARDGGQLPLLVVFSAMIFGTITNQDLPVIKCVLINHKNNDSS...</td>\n",
       "      <td>7QDP_2</td>\n",
       "      <td>P36888</td>\n",
       "      <td>14.400162</td>\n",
       "      <td>83</td>\n",
       "      <td>345</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Cc1ccccc1-c1c(-c2ccc3[nH]nc(N)c3c2)nnn1Cc1ccccc1</td>\n",
       "      <td>MPHPRRYHSSERGSRGSYREHYRSRKHKRRRSRSWSSSSDRTRRRR...</td>\n",
       "      <td>6KHE_1</td>\n",
       "      <td>P49760</td>\n",
       "      <td>12.399998</td>\n",
       "      <td>104</td>\n",
       "      <td>495</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CC(C)(C(N)=O)n1cc(-c2cnc(N)c3c(-c4ccc(NC(=O)Nc...</td>\n",
       "      <td>MSGRPRTTSFAESCKPVQQPSAFGSMKVSRDKDGSKVTTVVATPGQ...</td>\n",
       "      <td>1I09_1</td>\n",
       "      <td>P49841</td>\n",
       "      <td>11.400000</td>\n",
       "      <td>106</td>\n",
       "      <td>163</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Fc1ccc(-c2ccc3nccn3n2)cn1</td>\n",
       "      <td>MGAIGLLWLLPLLLSTAAVGSGMGTGQRAGSPAAGPPLQPREPLSY...</td>\n",
       "      <td>7NX3_1</td>\n",
       "      <td>Q9UM73</td>\n",
       "      <td>10.100000</td>\n",
       "      <td>222</td>\n",
       "      <td>1372</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>NCCCOc1cc2c(c(-c3ccc(Nc4nc5ccccc5o4)cc3)c1)CNC2=O</td>\n",
       "      <td>MASSSGSKAEFIVGGKYKLVRKIGSGSFGDIYLAINITNGEEVAVK...</td>\n",
       "      <td>5FQD_3</td>\n",
       "      <td>P48729</td>\n",
       "      <td>12.699998</td>\n",
       "      <td>97</td>\n",
       "      <td>254</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19704</th>\n",
       "      <td>CN(C)c1cc2sncc2cc1NC(=O)C(=O)O</td>\n",
       "      <td>MSELEEDFAKILMLKEERIKELEKRLSEKEEEIQELKRKLHKCQSV...</td>\n",
       "      <td>7LV3_1</td>\n",
       "      <td>Q13976</td>\n",
       "      <td>11.800001</td>\n",
       "      <td>159</td>\n",
       "      <td>467</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19705</th>\n",
       "      <td>CN1CCC1COc1cncc(CCc2ccncc2)c1</td>\n",
       "      <td>MAPFLRIAFNSYELGSLQAEDEANQPFCAVKMKEALSTERGKTLVQ...</td>\n",
       "      <td>1XJD_1</td>\n",
       "      <td>Q05655</td>\n",
       "      <td>11.200000</td>\n",
       "      <td>141</td>\n",
       "      <td>967</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19706</th>\n",
       "      <td>O=C(CO)N1CCC(c2[nH]nc(-c3ccc(Cl)cc3F)c2-c2ccnc...</td>\n",
       "      <td>MSSWIRWHGPAMARLWGFCWLVVGFWRAAFACPTSCKCSASRIWCS...</td>\n",
       "      <td>4ASZ_1</td>\n",
       "      <td>Q16620</td>\n",
       "      <td>11.400000</td>\n",
       "      <td>176</td>\n",
       "      <td>104</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19707</th>\n",
       "      <td>NNc1cc(N2CCOCC2)nc(OCCc2ccccn2)n1</td>\n",
       "      <td>MATCIGEKIEDFKVGNLLGKGSFAGVYRAESIHTGLEVAIKMIDKK...</td>\n",
       "      <td>3COK_1</td>\n",
       "      <td>O00444</td>\n",
       "      <td>11.500000</td>\n",
       "      <td>4</td>\n",
       "      <td>968</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19708</th>\n",
       "      <td>O=c1c(NCc2ccc(Cl)c(Cl)c2)c(Nc2ccncc2)c1=O</td>\n",
       "      <td>MADEDLIFRLEGVDGGQSPRAGHDGDSDGDSDDEEGYFICPITDDP...</td>\n",
       "      <td>7SHQ_1</td>\n",
       "      <td>O00418</td>\n",
       "      <td>11.500000</td>\n",
       "      <td>3</td>\n",
       "      <td>258</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>19681 rows × 7 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                     compound_iso_smiles  \\\n",
       "0            Cc1ccccc1NC(=O)Nc1ccc(-c2coc3ncnc(N)c23)cc1   \n",
       "1       Cc1ccccc1-c1c(-c2ccc3[nH]nc(N)c3c2)nnn1Cc1ccccc1   \n",
       "2      CC(C)(C(N)=O)n1cc(-c2cnc(N)c3c(-c4ccc(NC(=O)Nc...   \n",
       "3                              Fc1ccc(-c2ccc3nccn3n2)cn1   \n",
       "4      NCCCOc1cc2c(c(-c3ccc(Nc4nc5ccccc5o4)cc3)c1)CNC2=O   \n",
       "...                                                  ...   \n",
       "19704                     CN(C)c1cc2sncc2cc1NC(=O)C(=O)O   \n",
       "19705                      CN1CCC1COc1cncc(CCc2ccncc2)c1   \n",
       "19706  O=C(CO)N1CCC(c2[nH]nc(-c3ccc(Cl)cc3F)c2-c2ccnc...   \n",
       "19707                  NNc1cc(N2CCOCC2)nc(OCCc2ccccn2)n1   \n",
       "19708          O=c1c(NCc2ccc(Cl)c(Cl)c2)c(Nc2ccncc2)c1=O   \n",
       "\n",
       "                                         target_sequence target_pdb_id  \\\n",
       "0      MPALARDGGQLPLLVVFSAMIFGTITNQDLPVIKCVLINHKNNDSS...        7QDP_2   \n",
       "1      MPHPRRYHSSERGSRGSYREHYRSRKHKRRRSRSWSSSSDRTRRRR...        6KHE_1   \n",
       "2      MSGRPRTTSFAESCKPVQQPSAFGSMKVSRDKDGSKVTTVVATPGQ...        1I09_1   \n",
       "3      MGAIGLLWLLPLLLSTAAVGSGMGTGQRAGSPAAGPPLQPREPLSY...        7NX3_1   \n",
       "4      MASSSGSKAEFIVGGKYKLVRKIGSGSFGDIYLAINITNGEEVAVK...        5FQD_3   \n",
       "...                                                  ...           ...   \n",
       "19704  MSELEEDFAKILMLKEERIKELEKRLSEKEEEIQELKRKLHKCQSV...        7LV3_1   \n",
       "19705  MAPFLRIAFNSYELGSLQAEDEANQPFCAVKMKEALSTERGKTLVQ...        1XJD_1   \n",
       "19706  MSSWIRWHGPAMARLWGFCWLVVGFWRAAFACPTSCKCSASRIWCS...        4ASZ_1   \n",
       "19707  MATCIGEKIEDFKVGNLLGKGSFAGVYRAESIHTGLEVAIKMIDKK...        3COK_1   \n",
       "19708  MADEDLIFRLEGVDGGQSPRAGHDGDSDGDSDDEEGYFICPITDDP...        7SHQ_1   \n",
       "\n",
       "      target_original_key   affinity  protein_id  drug_id  \n",
       "0                  P36888  14.400162          83      345  \n",
       "1                  P49760  12.399998         104      495  \n",
       "2                  P49841  11.400000         106      163  \n",
       "3                  Q9UM73  10.100000         222     1372  \n",
       "4                  P48729  12.699998          97      254  \n",
       "...                   ...        ...         ...      ...  \n",
       "19704              Q13976  11.800001         159      467  \n",
       "19705              Q05655  11.200000         141      967  \n",
       "19706              Q16620  11.400000         176      104  \n",
       "19707              O00444  11.500000           4      968  \n",
       "19708              O00418  11.500000           3      258  \n",
       "\n",
       "[19681 rows x 7 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filter_by_tgt_ori_key(df, dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c2d771b6-9fad-49e8-ab22-0415af2f8b46",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|███████████████████████                                                                                            | 1/5 [00:01<00:04,  1.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 0 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|██████████████████████████████████████████████                                                                     | 2/5 [00:02<00:03,  1.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 1 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|█████████████████████████████████████████████████████████████████████                                              | 3/5 [00:03<00:02,  1.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 2 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████████████████████████████████████████████████████████████████████████████████████████                       | 4/5 [00:04<00:01,  1.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 3 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:05<00:00,  1.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis fold 4 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      " 20%|███████████████████████                                                                                            | 1/5 [00:03<00:15,  3.77s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 0 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|██████████████████████████████████████████████                                                                     | 2/5 [00:07<00:11,  3.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 1 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|█████████████████████████████████████████████████████████████████████                                              | 3/5 [00:11<00:07,  3.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 2 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████████████████████████████████████████████████████████████████████████████████████████                       | 4/5 [00:15<00:03,  3.79s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 3 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:18<00:00,  3.79s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kiba fold 4 clear!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "datasets = ['davis', 'kiba']\n",
    "folds = '01234'\n",
    "\n",
    "for dataset in datasets:\n",
    "    # test df\n",
    "    test_df = pd.read_csv(f'data/{dataset}/{dataset}_test.csv')\n",
    "    \n",
    "    filter_by_tgt_ori_key(test_df, dataset).to_csv(f'data/{dataset}/{dataset}_test_cm.csv', index=False)\n",
    "    \n",
    "    for fold in tqdm(folds):\n",
    "        # train & val df\n",
    "        train_df = pd.read_csv(f'data/{dataset}/{dataset}_train_fold_{fold}.csv')\n",
    "        valid_df = pd.read_csv(f'data/{dataset}/{dataset}_valid_fold_{fold}.csv')\n",
    "\n",
    "        filter_by_tgt_ori_key(train_df, dataset).to_csv(f'data/{dataset}/{dataset}_train_fold_{fold}_cm.csv', index=False)\n",
    "        filter_by_tgt_ori_key(valid_df, dataset).to_csv(f'data/{dataset}/{dataset}_valid_fold_{fold}_cm.csv', index=False)\n",
    "        \n",
    "        print(dataset, f'fold {fold} clear!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c73ff64c-2c29-446a-9087-fc800894e6f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "davis cold clear!\n",
      "kiba cold clear!\n"
     ]
    }
   ],
   "source": [
    "datasets = ['davis', 'kiba']\n",
    "\n",
    "for dataset in datasets:\n",
    "    # cold df\n",
    "    cold_df = pd.read_csv(f'data/{dataset}_cold.csv')\n",
    "\n",
    "    filter_by_tgt_ori_key(cold_df, dataset).to_csv(f'data/{dataset}_cold_cm.csv', index=False)\n",
    "    \n",
    "    print(dataset, f'cold clear!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "064a6455-f450-4836-9c9c-808acee94cb3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pl102",
   "language": "python",
   "name": "pl102"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}