Notebooks/human_random.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e55e5398-b351-402d-a535-567ad61aa896",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def load_dataset(sub_dataset, dataset_name):\n",
    "    base_path =  Path('Path_to_base_directory')\n",
    "\n",
    "    ############################# LOAD LLM Embeddigs ###############################################################\n",
    "    \n",
    "    llm_embeddings_path = base_path / f\"embeddings/LLM/{sub_dataset}\"\n",
    "    # drug\n",
    "    smile_names = np.load(llm_embeddings_path / f\"drug/{sub_dataset}_smiles.npy\", allow_pickle=True)\n",
    "    can_smile_names = np.load(llm_embeddings_path / f\"drug/{sub_dataset}_canonical_smiles.npy\", allow_pickle=True)\n",
    "    drug_embeddings = np.load(llm_embeddings_path / f\"drug/{sub_dataset}_molecule_embeddings.npy\", allow_pickle=True)\n",
    "    #gene\n",
    "    sequences_names = np.load(llm_embeddings_path / f\"target/{sub_dataset}_target_sequences.npy\", allow_pickle=True)\n",
    "    gene_embeddings = np.load(llm_embeddings_path / f\"target/{sub_dataset}_sequence_embeddings.npy\", allow_pickle=True)\n",
    "\n",
    "    \n",
    "    can_drugs_llm = pd.DataFrame({\n",
    "        'smiles': smile_names,\n",
    "        'can_smiles': can_smile_names,\n",
    "        'drug_llm_embeddings': drug_embeddings.tolist()  \n",
    "    })\n",
    "\n",
    "    drugs_llm = can_drugs_llm.dropna().reset_index(drop=True)\n",
    "\n",
    "    genes_llm = pd.DataFrame({\n",
    "        'sequences': sequences_names,\n",
    "        'protein_llm_embeddings': gene_embeddings.tolist()  \n",
    "    })\n",
    "\n",
    "    ############################# LOAD  Structural Features ###############################################################\n",
    "   \n",
    "    topological_embeddings_path = base_path / f\"embeddings/TDA/{dataset_name}\"\n",
    "    \n",
    "    # protein contact_map \n",
    "    protein_indexes = np.load(topological_embeddings_path / f\"02_protein_index.npy\")\n",
    "    protein_embeddings = np.load(topological_embeddings_path / f\"02_protein_contact_embeddings.npy\")\n",
    "    protein_sturct_embedding_df = pd.DataFrame({\"protein_index\":  protein_indexes, \"protein_struc_embeddings\": list(protein_embeddings)})\n",
    "\n",
    "\n",
    "    # Drug moleculer images\n",
    "    drug_indexes = np.load(topological_embeddings_path / f\"01_mol_names.npy\")\n",
    "    drug_struc_embeddings = np.load(topological_embeddings_path / f\"01_mol_image_embeddings.npy\")\n",
    "    drug_struct_embedding_df = pd.DataFrame({\"drug_index\": drug_indexes, \"drug_struc_embeddings\": list(drug_struc_embeddings)})\n",
    "\n",
    "\n",
    "    ############################# LOAD TRAIN VALIDATION AND TEST DATA ###############################################################   \n",
    "   \n",
    "    data_path = base_path / \"data\"\n",
    "    \n",
    "    task_paths = {\n",
    "        \"biosnap_random\": data_path / \"biosnap/random\",\n",
    "        \"human_random\": data_path / \"human/random\",\n",
    "        \"human_cold\": data_path / \"human/cold\"\n",
    "    }\n",
    "\n",
    "    # Validate dataset name\n",
    "    if sub_dataset not in task_paths:\n",
    "        raise ValueError(f\"Dataset name {sub_dataset} not available. Available datasets are {list(task_paths.keys())}.\")\n",
    "\n",
    "    dataset_path = task_paths[sub_dataset]\n",
    "    \n",
    "    train_file = dataset_path / 'train.csv'\n",
    "    valid_file = dataset_path / 'val.csv'\n",
    "    test_file = dataset_path / 'test.csv'\n",
    "\n",
    "    \n",
    "    train = pd.read_csv(train_file)\n",
    "    valid = pd.read_csv(valid_file)\n",
    "    test = pd.read_csv(test_file)\n",
    "\n",
    "    train['Set'] = \"Train\"\n",
    "    valid['Set'] = \"Valid\"\n",
    "    test['Set'] = \"Test\"\n",
    "\n",
    "    full_dataset = pd.concat([train, valid, test], ignore_index=True)\n",
    "\n",
    "    ############################# MERGE EMBEDDINGS ############################################################### \n",
    "    #Gene\n",
    "    genes_index = genes_llm.reset_index()\n",
    "    genes_index['index'] = genes_index['index'].astype(int)\n",
    "    protein_sturct_embedding_df['protein_index'] = protein_sturct_embedding_df['protein_index'].astype(int)\n",
    "    genes = pd.merge(genes_index,  protein_sturct_embedding_df, left_on = 'index', right_on ='protein_index')\n",
    "    genes.drop(columns=['index', 'protein_index'], inplace = True)\n",
    "    #Drug\n",
    "    drugs_index = drugs_llm.reset_index()\n",
    "    drugs_index['index'] = drugs_index['index'].astype(int)\n",
    "    drug_struct_embedding_df['drug_index'] = drug_struct_embedding_df['drug_index'].astype(int)\n",
    "    drugs = pd.merge(drugs_index, drug_struct_embedding_df, left_on = 'index', right_on ='drug_index')\n",
    "    drugs.drop(columns=['index', 'drug_index', 'can_smiles'], inplace = True)\n",
    "    \n",
    "\n",
    "    # Merge with drug and gene embeddings\n",
    "    drug_data = pd.merge(full_dataset, drugs, left_on='SMILES', right_on='smiles', how='inner')   \n",
    "    drug_data = drug_data.drop(columns=['smiles'])\n",
    "    \n",
    "    # Merge drug data with gene data\n",
    "    target_data = pd.merge(drug_data, genes, left_on='Protein', right_on='sequences', how='inner')\n",
    "    target_data = target_data.drop(columns=['sequences'])\n",
    "   \n",
    "\n",
    "\n",
    "    return drugs, genes, target_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "867a015f-ae00-4e10-902f-a50534ed262b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = 'human'\n",
    "sub_dataset ='human_random'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "71be2699-11ac-45e7-8b48-4f94ce35f774",
   "metadata": {},
   "outputs": [],
   "source": [
    "drugs, genes, df = load_dataset(sub_dataset, dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "782d06fb-256b-4541-bec3-4728813708f1",
   "metadata": {},
   "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>smiles</th>\n",
       "      <th>drug_llm_embeddings</th>\n",
       "      <th>drug_struc_embeddings</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>CC(C)[C@H](NS(=O)(=O)C1=CC=C(C=C1)C2=CC=C(C=C2...</td>\n",
       "      <td>[0.12668180465698242, 0.7376371622085571, 0.43...</td>\n",
       "      <td>[49.0, 13.0, 11.0, 9.0, 9.0, 9.0, 11.0, 11.0, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2...</td>\n",
       "      <td>[0.354166716337204, -1.0170406103134155, 0.700...</td>\n",
       "      <td>[1.0, 2.0, 4.0, 4.0, 5.0, 6.0, 11.0, 11.0, 16....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CC12CCC3C(C1CCC2=O)CCC4=CC(=O)CCC34CO</td>\n",
       "      <td>[0.8736435174942017, -0.547260582447052, -0.27...</td>\n",
       "      <td>[64.0, 4.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>C1CN(CCC1(C2=CC=C(C=C2)Cl)O)CCCC(=O)C3=CC=C(C=...</td>\n",
       "      <td>[0.5970272421836853, 0.024165719747543335, 0.5...</td>\n",
       "      <td>[55.0, 11.0, 10.0, 9.0, 9.0, 8.0, 7.0, 9.0, 9....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>C1C[C@@H]2CN[C@H](C[C@@H]2C[C@@H]1CCC3=NNN=N3)...</td>\n",
       "      <td>[-0.04980764165520668, 0.36217665672302246, 0....</td>\n",
       "      <td>[49.0, 9.0, 6.0, 5.0, 6.0, 7.0, 17.0, 20.0, 22...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2721</th>\n",
       "      <td>C1=CC=C(C=C1)CN[C@@H](C2=CC=CC=C2)P(=O)(O)O</td>\n",
       "      <td>[0.639024019241333, 0.46341198682785034, -0.24...</td>\n",
       "      <td>[62.0, 7.0, 7.0, 7.0, 7.0, 8.0, 14.0, 15.0, 14...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2722</th>\n",
       "      <td>C1CCC(CC1)COC2=NC(=NC3=C2NC=N3)NC4=CC(=CC=C4)S...</td>\n",
       "      <td>[0.25831153988838196, 0.45374321937561035, -0....</td>\n",
       "      <td>[93.0, 16.0, 14.0, 13.0, 12.0, 11.0, 40.0, 44....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2723</th>\n",
       "      <td>CC1=C(C=C(C=C1)OC)OCC2=CC=NC=C2</td>\n",
       "      <td>[-0.033933546394109726, 0.9028526544570923, 0....</td>\n",
       "      <td>[68.0, 12.0, 11.0, 9.0, 8.0, 8.0, 15.0, 13.0, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2724</th>\n",
       "      <td>C1CC[N+]2(C1)[C@@H]3CC[C@H]2CC(C3)OC(=O)C(C4=C...</td>\n",
       "      <td>[0.6405504941940308, 0.21705366671085358, 0.36...</td>\n",
       "      <td>[85.0, 14.0, 12.0, 12.0, 12.0, 13.0, 17.0, 17....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2725</th>\n",
       "      <td>B1(OCCO1)OCCCC[NH3+]</td>\n",
       "      <td>[0.4124879539012909, -0.5095826387405396, 0.39...</td>\n",
       "      <td>[26.0, 6.0, 6.0, 6.0, 6.0, 6.0, 8.0, 9.0, 11.0...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2726 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 smiles  \\\n",
       "0     CC(C)[C@H](NS(=O)(=O)C1=CC=C(C=C1)C2=CC=C(C=C2...   \n",
       "1     [O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2...   \n",
       "2                 CC12CCC3C(C1CCC2=O)CCC4=CC(=O)CCC34CO   \n",
       "3     C1CN(CCC1(C2=CC=C(C=C2)Cl)O)CCCC(=O)C3=CC=C(C=...   \n",
       "4     C1C[C@@H]2CN[C@H](C[C@@H]2C[C@@H]1CCC3=NNN=N3)...   \n",
       "...                                                 ...   \n",
       "2721        C1=CC=C(C=C1)CN[C@@H](C2=CC=CC=C2)P(=O)(O)O   \n",
       "2722  C1CCC(CC1)COC2=NC(=NC3=C2NC=N3)NC4=CC(=CC=C4)S...   \n",
       "2723                    CC1=C(C=C(C=C1)OC)OCC2=CC=NC=C2   \n",
       "2724  C1CC[N+]2(C1)[C@@H]3CC[C@H]2CC(C3)OC(=O)C(C4=C...   \n",
       "2725                               B1(OCCO1)OCCCC[NH3+]   \n",
       "\n",
       "                                    drug_llm_embeddings  \\\n",
       "0     [0.12668180465698242, 0.7376371622085571, 0.43...   \n",
       "1     [0.354166716337204, -1.0170406103134155, 0.700...   \n",
       "2     [0.8736435174942017, -0.547260582447052, -0.27...   \n",
       "3     [0.5970272421836853, 0.024165719747543335, 0.5...   \n",
       "4     [-0.04980764165520668, 0.36217665672302246, 0....   \n",
       "...                                                 ...   \n",
       "2721  [0.639024019241333, 0.46341198682785034, -0.24...   \n",
       "2722  [0.25831153988838196, 0.45374321937561035, -0....   \n",
       "2723  [-0.033933546394109726, 0.9028526544570923, 0....   \n",
       "2724  [0.6405504941940308, 0.21705366671085358, 0.36...   \n",
       "2725  [0.4124879539012909, -0.5095826387405396, 0.39...   \n",
       "\n",
       "                                  drug_struc_embeddings  \n",
       "0     [49.0, 13.0, 11.0, 9.0, 9.0, 9.0, 11.0, 11.0, ...  \n",
       "1     [1.0, 2.0, 4.0, 4.0, 5.0, 6.0, 11.0, 11.0, 16....  \n",
       "2     [64.0, 4.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,...  \n",
       "3     [55.0, 11.0, 10.0, 9.0, 9.0, 8.0, 7.0, 9.0, 9....  \n",
       "4     [49.0, 9.0, 6.0, 5.0, 6.0, 7.0, 17.0, 20.0, 22...  \n",
       "...                                                 ...  \n",
       "2721  [62.0, 7.0, 7.0, 7.0, 7.0, 8.0, 14.0, 15.0, 14...  \n",
       "2722  [93.0, 16.0, 14.0, 13.0, 12.0, 11.0, 40.0, 44....  \n",
       "2723  [68.0, 12.0, 11.0, 9.0, 8.0, 8.0, 15.0, 13.0, ...  \n",
       "2724  [85.0, 14.0, 12.0, 12.0, 12.0, 13.0, 17.0, 17....  \n",
       "2725  [26.0, 6.0, 6.0, 6.0, 6.0, 6.0, 8.0, 9.0, 11.0...  \n",
       "\n",
       "[2726 rows x 3 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drugs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "02105eea-efe6-45cd-a201-c503dbaa4ef5",
   "metadata": {},
   "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>sequences</th>\n",
       "      <th>protein_llm_embeddings</th>\n",
       "      <th>protein_struc_embeddings</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>MILLTFSTGRRLDFVHHSGVFFLQTLLWILCATVCGTEQYFNVEVW...</td>\n",
       "      <td>[0.04457112401723862, 0.06250955909490585, 0.0...</td>\n",
       "      <td>[62.0, 62.0, 813.0, 375.0, 375.0, 250.0, 250.0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>MLPSASRERPGYRAGVAAPDLLDPKSAAQNSKPRLSFSTKPTVLAS...</td>\n",
       "      <td>[0.08137195557355881, 0.10094469040632248, 0.0...</td>\n",
       "      <td>[26.0, 26.0, 962.0, 354.0, 354.0, 230.0, 230.0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>MWSWKCLLFWAVLVTATLCTARPSPTLPEQAQPWGAPVEVESFLVH...</td>\n",
       "      <td>[0.052909743040800095, 0.09419899433851242, 0....</td>\n",
       "      <td>[88.0, 308.0, 171.0, 93.0, 24.0, 23.0, 9.0, 16...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>MPVRRGHVAPQNTFLDTIIRKFEGQSRKFIIANARVENCAVIYCND...</td>\n",
       "      <td>[0.034286316484212875, 0.07540300488471985, 0....</td>\n",
       "      <td>[79.0, 2015.0, 246.0, 246.0, 98.0, 98.0, 61.0,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>MAEDGEEAEFHFAALYISGQWPRLRADTDLQRLGSSAMAPSRKFFV...</td>\n",
       "      <td>[0.0838751345872879, 0.039654411375522614, 0.0...</td>\n",
       "      <td>[56.0, 56.0, 524.0, 524.0, 324.0, 324.0, 324.0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1996</th>\n",
       "      <td>MAQWEMLQNLDSPFQDQLHQLYSHSLLPVDIRQYLAVWIEDQNWQE...</td>\n",
       "      <td>[-0.03957125544548035, 0.02081611379981041, 0....</td>\n",
       "      <td>[40.0, 40.0, 1090.0, 1090.0, 412.0, 412.0, 412...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1997</th>\n",
       "      <td>MGLQACLLGLFALILSGKCSYSPEPDQRRTLPPGWVSLGRADPEEE...</td>\n",
       "      <td>[-0.013600670732557774, 0.03684832900762558, 0...</td>\n",
       "      <td>[70.0, 70.0, 1387.0, 499.0, 499.0, 499.0, 292....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1998</th>\n",
       "      <td>MRYRLAWLLHPALPSTFRSVLGARLPPPERLCGFQKKTYSKMNNPA...</td>\n",
       "      <td>[0.035878922790288925, 0.07692313194274902, 0....</td>\n",
       "      <td>[75.0, 75.0, 75.0, 2717.0, 2717.0, 399.0, 399....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1999</th>\n",
       "      <td>MKALCLLLLPVLGLLVSSKTLCSMEEAINERIQEVAGSLIFRAISS...</td>\n",
       "      <td>[0.06332018971443176, 0.01541752927005291, 0.0...</td>\n",
       "      <td>[126.0, 126.0, 28.0, 28.0, 46.0, 46.0, 46.0, 5...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2000</th>\n",
       "      <td>MSQEGDYGRWTISSSDESEEEKPKPDKPSTSSLLCARQGAANEPRY...</td>\n",
       "      <td>[-0.008833632804453373, 0.05300190672278404, 0...</td>\n",
       "      <td>[37.0, 37.0, 1009.0, 1009.0, 364.0, 364.0, 364...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2001 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              sequences  \\\n",
       "0     MILLTFSTGRRLDFVHHSGVFFLQTLLWILCATVCGTEQYFNVEVW...   \n",
       "1     MLPSASRERPGYRAGVAAPDLLDPKSAAQNSKPRLSFSTKPTVLAS...   \n",
       "2     MWSWKCLLFWAVLVTATLCTARPSPTLPEQAQPWGAPVEVESFLVH...   \n",
       "3     MPVRRGHVAPQNTFLDTIIRKFEGQSRKFIIANARVENCAVIYCND...   \n",
       "4     MAEDGEEAEFHFAALYISGQWPRLRADTDLQRLGSSAMAPSRKFFV...   \n",
       "...                                                 ...   \n",
       "1996  MAQWEMLQNLDSPFQDQLHQLYSHSLLPVDIRQYLAVWIEDQNWQE...   \n",
       "1997  MGLQACLLGLFALILSGKCSYSPEPDQRRTLPPGWVSLGRADPEEE...   \n",
       "1998  MRYRLAWLLHPALPSTFRSVLGARLPPPERLCGFQKKTYSKMNNPA...   \n",
       "1999  MKALCLLLLPVLGLLVSSKTLCSMEEAINERIQEVAGSLIFRAISS...   \n",
       "2000  MSQEGDYGRWTISSSDESEEEKPKPDKPSTSSLLCARQGAANEPRY...   \n",
       "\n",
       "                                 protein_llm_embeddings  \\\n",
       "0     [0.04457112401723862, 0.06250955909490585, 0.0...   \n",
       "1     [0.08137195557355881, 0.10094469040632248, 0.0...   \n",
       "2     [0.052909743040800095, 0.09419899433851242, 0....   \n",
       "3     [0.034286316484212875, 0.07540300488471985, 0....   \n",
       "4     [0.0838751345872879, 0.039654411375522614, 0.0...   \n",
       "...                                                 ...   \n",
       "1996  [-0.03957125544548035, 0.02081611379981041, 0....   \n",
       "1997  [-0.013600670732557774, 0.03684832900762558, 0...   \n",
       "1998  [0.035878922790288925, 0.07692313194274902, 0....   \n",
       "1999  [0.06332018971443176, 0.01541752927005291, 0.0...   \n",
       "2000  [-0.008833632804453373, 0.05300190672278404, 0...   \n",
       "\n",
       "                               protein_struc_embeddings  \n",
       "0     [62.0, 62.0, 813.0, 375.0, 375.0, 250.0, 250.0...  \n",
       "1     [26.0, 26.0, 962.0, 354.0, 354.0, 230.0, 230.0...  \n",
       "2     [88.0, 308.0, 171.0, 93.0, 24.0, 23.0, 9.0, 16...  \n",
       "3     [79.0, 2015.0, 246.0, 246.0, 98.0, 98.0, 61.0,...  \n",
       "4     [56.0, 56.0, 524.0, 524.0, 324.0, 324.0, 324.0...  \n",
       "...                                                 ...  \n",
       "1996  [40.0, 40.0, 1090.0, 1090.0, 412.0, 412.0, 412...  \n",
       "1997  [70.0, 70.0, 1387.0, 499.0, 499.0, 499.0, 292....  \n",
       "1998  [75.0, 75.0, 75.0, 2717.0, 2717.0, 399.0, 399....  \n",
       "1999  [126.0, 126.0, 28.0, 28.0, 46.0, 46.0, 46.0, 5...  \n",
       "2000  [37.0, 37.0, 1009.0, 1009.0, 364.0, 364.0, 364...  \n",
       "\n",
       "[2001 rows x 3 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c03b91b0-2890-46e9-bac9-d757c0bd044d",
   "metadata": {},
   "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>SMILES</th>\n",
       "      <th>Protein</th>\n",
       "      <th>Y</th>\n",
       "      <th>Set</th>\n",
       "      <th>drug_llm_embeddings</th>\n",
       "      <th>drug_struc_embeddings</th>\n",
       "      <th>protein_llm_embeddings</th>\n",
       "      <th>protein_struc_embeddings</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>CC(C)[C@H](NS(=O)(=O)C1=CC=C(C=C1)C2=CC=C(C=C2...</td>\n",
       "      <td>MILLTFSTGRRLDFVHHSGVFFLQTLLWILCATVCGTEQYFNVEVW...</td>\n",
       "      <td>1</td>\n",
       "      <td>Train</td>\n",
       "      <td>[0.12668180465698242, 0.7376371622085571, 0.43...</td>\n",
       "      <td>[49.0, 13.0, 11.0, 9.0, 9.0, 9.0, 11.0, 11.0, ...</td>\n",
       "      <td>[0.04457112401723862, 0.06250955909490585, 0.0...</td>\n",
       "      <td>[62.0, 62.0, 813.0, 375.0, 375.0, 250.0, 250.0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2...</td>\n",
       "      <td>MLPSASRERPGYRAGVAAPDLLDPKSAAQNSKPRLSFSTKPTVLAS...</td>\n",
       "      <td>0</td>\n",
       "      <td>Train</td>\n",
       "      <td>[0.354166716337204, -1.0170406103134155, 0.700...</td>\n",
       "      <td>[1.0, 2.0, 4.0, 4.0, 5.0, 6.0, 11.0, 11.0, 16....</td>\n",
       "      <td>[0.08137195557355881, 0.10094469040632248, 0.0...</td>\n",
       "      <td>[26.0, 26.0, 962.0, 354.0, 354.0, 230.0, 230.0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CC12CCC3C(C1CCC2=O)CCC4=CC(=O)CCC34CO</td>\n",
       "      <td>MWSWKCLLFWAVLVTATLCTARPSPTLPEQAQPWGAPVEVESFLVH...</td>\n",
       "      <td>0</td>\n",
       "      <td>Train</td>\n",
       "      <td>[0.8736435174942017, -0.547260582447052, -0.27...</td>\n",
       "      <td>[64.0, 4.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,...</td>\n",
       "      <td>[0.052909743040800095, 0.09419899433851242, 0....</td>\n",
       "      <td>[88.0, 308.0, 171.0, 93.0, 24.0, 23.0, 9.0, 16...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>C1CN(CCC1(C2=CC=C(C=C2)Cl)O)CCCC(=O)C3=CC=C(C=...</td>\n",
       "      <td>MPVRRGHVAPQNTFLDTIIRKFEGQSRKFIIANARVENCAVIYCND...</td>\n",
       "      <td>1</td>\n",
       "      <td>Train</td>\n",
       "      <td>[0.5970272421836853, 0.024165719747543335, 0.5...</td>\n",
       "      <td>[55.0, 11.0, 10.0, 9.0, 9.0, 8.0, 7.0, 9.0, 9....</td>\n",
       "      <td>[0.034286316484212875, 0.07540300488471985, 0....</td>\n",
       "      <td>[79.0, 2015.0, 246.0, 246.0, 98.0, 98.0, 61.0,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>C1C[C@@H]2CN[C@H](C[C@@H]2C[C@@H]1CCC3=NNN=N3)...</td>\n",
       "      <td>MAEDGEEAEFHFAALYISGQWPRLRADTDLQRLGSSAMAPSRKFFV...</td>\n",
       "      <td>0</td>\n",
       "      <td>Train</td>\n",
       "      <td>[-0.04980764165520668, 0.36217665672302246, 0....</td>\n",
       "      <td>[49.0, 9.0, 6.0, 5.0, 6.0, 7.0, 17.0, 20.0, 22...</td>\n",
       "      <td>[0.0838751345872879, 0.039654411375522614, 0.0...</td>\n",
       "      <td>[56.0, 56.0, 524.0, 524.0, 324.0, 324.0, 324.0...</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",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5992</th>\n",
       "      <td>C1=NC2=C(C(=N1)N)N=CN2[C@H]3[C@@H]([C@@H]([C@H...</td>\n",
       "      <td>MSANATLKPLCPILEQMSRLQSHSNTSIRYIDHAAVLLHGLASLLG...</td>\n",
       "      <td>1</td>\n",
       "      <td>Test</td>\n",
       "      <td>[0.8082612156867981, 0.481841504573822, 0.2084...</td>\n",
       "      <td>[27.0, 13.0, 12.0, 12.0, 11.0, 11.0, 24.0, 22....</td>\n",
       "      <td>[-0.05475754290819168, 0.05181394889950752, 0....</td>\n",
       "      <td>[58.0, 58.0, 1300.0, 366.0, 366.0, 172.0, 172....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5993</th>\n",
       "      <td>C[C@H]1[C@H]([C@H](CC(O1)O[C@H]2C[C@@](CC3=C(C...</td>\n",
       "      <td>MGAIGLLWLLPLLLSTAAVGSGMGTGQRAGSPAAGPPLQPREPLSY...</td>\n",
       "      <td>0</td>\n",
       "      <td>Test</td>\n",
       "      <td>[0.7286798357963562, 0.10095473378896713, 0.35...</td>\n",
       "      <td>[80.0, 17.0, 11.0, 10.0, 12.0, 13.0, 16.0, 14....</td>\n",
       "      <td>[0.024710068479180336, 0.03366069868206978, 0....</td>\n",
       "      <td>[167.0, 167.0, 167.0, 167.0, 167.0, 167.0, 412...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5994</th>\n",
       "      <td>B1(OCCO1)OCCCC[NH3+]</td>\n",
       "      <td>MATAGNPWGWFLGYLILGVAGSLVSGSCSQIINGEDCSPHSQPWQA...</td>\n",
       "      <td>1</td>\n",
       "      <td>Test</td>\n",
       "      <td>[0.4124879539012909, -0.5095826387405396, 0.39...</td>\n",
       "      <td>[26.0, 6.0, 6.0, 6.0, 6.0, 6.0, 8.0, 9.0, 11.0...</td>\n",
       "      <td>[0.001806111540645361, 0.04118546470999718, 0....</td>\n",
       "      <td>[110.0, 476.0, 324.0, 324.0, 88.0, 88.0, 42.0,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5995</th>\n",
       "      <td>CC1=C(C(CCC1)(C)C)/C=C/C(=C/C=C/C(=C/C=C/C=C(\\...</td>\n",
       "      <td>MATPPKRSCPSFSASSEGTRIKKISIEGNIAAGKSTFVNILKQLCE...</td>\n",
       "      <td>0</td>\n",
       "      <td>Test</td>\n",
       "      <td>[0.8428023457527161, 0.08140375465154648, 0.68...</td>\n",
       "      <td>[62.0, 10.0, 5.0, 4.0, 1.0, 1.0, 0.0, 0.0, 0.0...</td>\n",
       "      <td>[0.040746789425611496, 0.0875672921538353, 0.0...</td>\n",
       "      <td>[78.0, 766.0, 290.0, 290.0, 90.0, 90.0, 50.0, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5996</th>\n",
       "      <td>CC(C)[C@@H](C(=O)N)NC</td>\n",
       "      <td>MEQTEVLKPRTLADLIRILHQLFAGDEVNVEEVQAIMEAYESDPTE...</td>\n",
       "      <td>0</td>\n",
       "      <td>Test</td>\n",
       "      <td>[-0.8276767730712891, 0.37549686431884766, 0.2...</td>\n",
       "      <td>[37.0, 3.0, 3.0, 3.0, 2.0, 1.0, 12.0, 14.0, 16...</td>\n",
       "      <td>[0.07461965084075928, 0.051294729113578796, -0...</td>\n",
       "      <td>[128.0, 128.0, 350.0, 350.0, 190.0, 190.0, 190...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5997 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 SMILES  \\\n",
       "0     CC(C)[C@H](NS(=O)(=O)C1=CC=C(C=C1)C2=CC=C(C=C2...   \n",
       "1     [O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2].[O-2...   \n",
       "2                 CC12CCC3C(C1CCC2=O)CCC4=CC(=O)CCC34CO   \n",
       "3     C1CN(CCC1(C2=CC=C(C=C2)Cl)O)CCCC(=O)C3=CC=C(C=...   \n",
       "4     C1C[C@@H]2CN[C@H](C[C@@H]2C[C@@H]1CCC3=NNN=N3)...   \n",
       "...                                                 ...   \n",
       "5992  C1=NC2=C(C(=N1)N)N=CN2[C@H]3[C@@H]([C@@H]([C@H...   \n",
       "5993  C[C@H]1[C@H]([C@H](CC(O1)O[C@H]2C[C@@](CC3=C(C...   \n",
       "5994                               B1(OCCO1)OCCCC[NH3+]   \n",
       "5995  CC1=C(C(CCC1)(C)C)/C=C/C(=C/C=C/C(=C/C=C/C=C(\\...   \n",
       "5996                              CC(C)[C@@H](C(=O)N)NC   \n",
       "\n",
       "                                                Protein  Y    Set  \\\n",
       "0     MILLTFSTGRRLDFVHHSGVFFLQTLLWILCATVCGTEQYFNVEVW...  1  Train   \n",
       "1     MLPSASRERPGYRAGVAAPDLLDPKSAAQNSKPRLSFSTKPTVLAS...  0  Train   \n",
       "2     MWSWKCLLFWAVLVTATLCTARPSPTLPEQAQPWGAPVEVESFLVH...  0  Train   \n",
       "3     MPVRRGHVAPQNTFLDTIIRKFEGQSRKFIIANARVENCAVIYCND...  1  Train   \n",
       "4     MAEDGEEAEFHFAALYISGQWPRLRADTDLQRLGSSAMAPSRKFFV...  0  Train   \n",
       "...                                                 ... ..    ...   \n",
       "5992  MSANATLKPLCPILEQMSRLQSHSNTSIRYIDHAAVLLHGLASLLG...  1   Test   \n",
       "5993  MGAIGLLWLLPLLLSTAAVGSGMGTGQRAGSPAAGPPLQPREPLSY...  0   Test   \n",
       "5994  MATAGNPWGWFLGYLILGVAGSLVSGSCSQIINGEDCSPHSQPWQA...  1   Test   \n",
       "5995  MATPPKRSCPSFSASSEGTRIKKISIEGNIAAGKSTFVNILKQLCE...  0   Test   \n",
       "5996  MEQTEVLKPRTLADLIRILHQLFAGDEVNVEEVQAIMEAYESDPTE...  0   Test   \n",
       "\n",
       "                                    drug_llm_embeddings  \\\n",
       "0     [0.12668180465698242, 0.7376371622085571, 0.43...   \n",
       "1     [0.354166716337204, -1.0170406103134155, 0.700...   \n",
       "2     [0.8736435174942017, -0.547260582447052, -0.27...   \n",
       "3     [0.5970272421836853, 0.024165719747543335, 0.5...   \n",
       "4     [-0.04980764165520668, 0.36217665672302246, 0....   \n",
       "...                                                 ...   \n",
       "5992  [0.8082612156867981, 0.481841504573822, 0.2084...   \n",
       "5993  [0.7286798357963562, 0.10095473378896713, 0.35...   \n",
       "5994  [0.4124879539012909, -0.5095826387405396, 0.39...   \n",
       "5995  [0.8428023457527161, 0.08140375465154648, 0.68...   \n",
       "5996  [-0.8276767730712891, 0.37549686431884766, 0.2...   \n",
       "\n",
       "                                  drug_struc_embeddings  \\\n",
       "0     [49.0, 13.0, 11.0, 9.0, 9.0, 9.0, 11.0, 11.0, ...   \n",
       "1     [1.0, 2.0, 4.0, 4.0, 5.0, 6.0, 11.0, 11.0, 16....   \n",
       "2     [64.0, 4.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,...   \n",
       "3     [55.0, 11.0, 10.0, 9.0, 9.0, 8.0, 7.0, 9.0, 9....   \n",
       "4     [49.0, 9.0, 6.0, 5.0, 6.0, 7.0, 17.0, 20.0, 22...   \n",
       "...                                                 ...   \n",
       "5992  [27.0, 13.0, 12.0, 12.0, 11.0, 11.0, 24.0, 22....   \n",
       "5993  [80.0, 17.0, 11.0, 10.0, 12.0, 13.0, 16.0, 14....   \n",
       "5994  [26.0, 6.0, 6.0, 6.0, 6.0, 6.0, 8.0, 9.0, 11.0...   \n",
       "5995  [62.0, 10.0, 5.0, 4.0, 1.0, 1.0, 0.0, 0.0, 0.0...   \n",
       "5996  [37.0, 3.0, 3.0, 3.0, 2.0, 1.0, 12.0, 14.0, 16...   \n",
       "\n",
       "                                 protein_llm_embeddings  \\\n",
       "0     [0.04457112401723862, 0.06250955909490585, 0.0...   \n",
       "1     [0.08137195557355881, 0.10094469040632248, 0.0...   \n",
       "2     [0.052909743040800095, 0.09419899433851242, 0....   \n",
       "3     [0.034286316484212875, 0.07540300488471985, 0....   \n",
       "4     [0.0838751345872879, 0.039654411375522614, 0.0...   \n",
       "...                                                 ...   \n",
       "5992  [-0.05475754290819168, 0.05181394889950752, 0....   \n",
       "5993  [0.024710068479180336, 0.03366069868206978, 0....   \n",
       "5994  [0.001806111540645361, 0.04118546470999718, 0....   \n",
       "5995  [0.040746789425611496, 0.0875672921538353, 0.0...   \n",
       "5996  [0.07461965084075928, 0.051294729113578796, -0...   \n",
       "\n",
       "                               protein_struc_embeddings  \n",
       "0     [62.0, 62.0, 813.0, 375.0, 375.0, 250.0, 250.0...  \n",
       "1     [26.0, 26.0, 962.0, 354.0, 354.0, 230.0, 230.0...  \n",
       "2     [88.0, 308.0, 171.0, 93.0, 24.0, 23.0, 9.0, 16...  \n",
       "3     [79.0, 2015.0, 246.0, 246.0, 98.0, 98.0, 61.0,...  \n",
       "4     [56.0, 56.0, 524.0, 524.0, 324.0, 324.0, 324.0...  \n",
       "...                                                 ...  \n",
       "5992  [58.0, 58.0, 1300.0, 366.0, 366.0, 172.0, 172....  \n",
       "5993  [167.0, 167.0, 167.0, 167.0, 167.0, 167.0, 412...  \n",
       "5994  [110.0, 476.0, 324.0, 324.0, 88.0, 88.0, 42.0,...  \n",
       "5995  [78.0, 766.0, 290.0, 290.0, 90.0, 90.0, 50.0, ...  \n",
       "5996  [128.0, 128.0, 350.0, 350.0, 190.0, 190.0, 190...  \n",
       "\n",
       "[5997 rows x 8 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "371541c0-3a12-4003-b255-eccddd94edbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SMILES                      0\n",
       "Protein                     0\n",
       "Y                           0\n",
       "Set                         0\n",
       "drug_llm_embeddings         0\n",
       "drug_struc_embeddings       0\n",
       "protein_llm_embeddings      0\n",
       "protein_struc_embeddings    0\n",
       "dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.isna().sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "efced911-4f70-466d-a63b-54cd707a524b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "from torch_geometric.data import HeteroData\n",
    "import torch_geometric.transforms as T\n",
    "import pandas as pd\n",
    "import random\n",
    "random.seed(42)\n",
    "\n",
    "# unique drugs and genes\n",
    "unique_drug_id = drugs['smiles'].unique()\n",
    "unique_gene_id = genes['sequences'].unique()\n",
    "\n",
    "# their ids\n",
    "drug_id_map = {drug: idx for idx, drug in enumerate(unique_drug_id)}\n",
    "gene_id_map = {gene: idx for idx, gene in enumerate(unique_gene_id)}\n",
    "\n",
    "# and embeddings \n",
    "drug_llm_embeddings = np.array([np.array(emb) for emb in drugs['drug_llm_embeddings']])\n",
    "drug_struct_embeddings = np.array([np.array(emb) for emb in drugs['drug_struc_embeddings']])\n",
    "gene_llm_embeddings = np.array([np.array(emb) for emb in genes['protein_llm_embeddings']])\n",
    "protein_struct_embeddings = np.array([np.array(emb) for emb in genes['protein_struc_embeddings']])\n",
    "\n",
    "\n",
    "data = HeteroData()\n",
    "\n",
    "data[\"drug\"].node_id = torch.arange(len(unique_drug_id))\n",
    "data[\"gene\"].node_id = torch.arange(len(unique_gene_id))\n",
    "\n",
    "# node features\n",
    "data[\"drug\"].xl = torch.tensor(drug_llm_embeddings, dtype=torch.float)\n",
    "data[\"drug\"].xs = torch.tensor(drug_struct_embeddings, dtype=torch.float)\n",
    "data[\"gene\"].xl = torch.tensor(gene_llm_embeddings, dtype=torch.float)\n",
    "data[\"gene\"].xs = torch.tensor(protein_struct_embeddings, dtype=torch.float)\n",
    "\n",
    "\n",
    "\n",
    "edge_indices, edge_labels, edge_splits = [], [], []\n",
    "\n",
    "for _, row in df.iterrows():\n",
    "    drug_idx = drug_id_map[row[\"SMILES\"]]\n",
    "    gene_idx = gene_id_map[row[\"Protein\"]]\n",
    "    edge_indices.append([drug_idx, gene_idx])\n",
    "    edge_labels.append(row[\"Y\"])\n",
    "    edge_splits.append(row[\"Set\"])\n",
    "\n",
    "# create tensors\n",
    "edge_indices = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()\n",
    "edge_labels = torch.tensor(edge_labels, dtype=torch.float)\n",
    "\n",
    "# and add them to data\n",
    "data[\"drug\", \"interacts_with\", \"gene\"].edge_index = edge_indices\n",
    "data[\"drug\", \"interacts_with\", \"gene\"].edge_label = edge_labels\n",
    "\n",
    "# Split data according to Set Labels\n",
    "edge_splits = np.array(edge_splits)\n",
    "train_mask = torch.tensor(edge_splits == \"Train\", dtype=torch.bool)\n",
    "valid_mask = torch.tensor(edge_splits == \"Valid\", dtype=torch.bool)\n",
    "test_mask = torch.tensor(edge_splits == \"Test\", dtype=torch.bool)\n",
    "\n",
    "data[\"drug\", \"interacts_with\", \"gene\"].train_mask = train_mask\n",
    "data[\"drug\", \"interacts_with\", \"gene\"].valid_mask = valid_mask\n",
    "data[\"drug\", \"interacts_with\", \"gene\"].test_mask = test_mask\n",
    "data = T.ToUndirected()(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "64012972-ecac-4afc-baaf-a005cb79b260",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  drug={\n",
       "    node_id=[2726],\n",
       "    xl=[2726, 768],\n",
       "    xs=[2726, 1200],\n",
       "  },\n",
       "  gene={\n",
       "    node_id=[2001],\n",
       "    xl=[2001, 1024],\n",
       "    xs=[2001, 1200],\n",
       "  },\n",
       "  (drug, interacts_with, gene)={\n",
       "    edge_index=[2, 5997],\n",
       "    edge_label=[5997],\n",
       "    train_mask=[5997],\n",
       "    valid_mask=[5997],\n",
       "    test_mask=[5997],\n",
       "  },\n",
       "  (gene, rev_interacts_with, drug)={\n",
       "    edge_index=[2, 5997],\n",
       "    edge_label=[5997],\n",
       "    train_mask=[5997],\n",
       "    valid_mask=[5997],\n",
       "    test_mask=[5997],\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4ce529fe-f6b9-47f0-806c-38e51b0923ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===== Full Dataset =====\n",
      "Number of drug nodes: 2726\n",
      "Number of gene nodes: 2001\n",
      "Number of edges: 5997\n",
      "\n",
      "===== Training Dataset =====\n",
      "Number of drug nodes: 2726\n",
      "Number of gene nodes: 2001\n",
      "Number of edges: 4197\n",
      "\n",
      "===== Validation Dataset =====\n",
      "Number of drug nodes: 2726\n",
      "Number of gene nodes: 2001\n",
      "Number of edges: 600\n",
      "\n",
      "===== Test Dataset =====\n",
      "Number of drug nodes: 2726\n",
      "Number of gene nodes: 2001\n",
      "Number of edges: 1200\n"
     ]
    }
   ],
   "source": [
    "# Full dataset\n",
    "print(\"===== Full Dataset =====\")\n",
    "print(\"Number of drug nodes:\", data['drug'].num_nodes)\n",
    "print(\"Number of gene nodes:\", data['gene'].num_nodes)\n",
    "print(\"Number of edges:\", data['drug', 'interacts_with', 'gene'].num_edges)\n",
    "\n",
    "# Train\n",
    "train_mask = data[\"drug\", \"interacts_with\", \"gene\"].train_mask\n",
    "print(\"\\n===== Training Dataset =====\")\n",
    "print(\"Number of drug nodes:\", data['drug'].num_nodes)\n",
    "print(\"Number of gene nodes:\", data['gene'].num_nodes)\n",
    "print(\"Number of edges:\", train_mask.sum().item())\n",
    "\n",
    "# Validation\n",
    "valid_mask = data[\"drug\", \"interacts_with\", \"gene\"].valid_mask\n",
    "print(\"\\n===== Validation Dataset =====\")\n",
    "print(\"Number of drug nodes:\", data['drug'].num_nodes)\n",
    "print(\"Number of gene nodes:\", data['gene'].num_nodes)\n",
    "print(\"Number of edges:\", valid_mask.sum().item())\n",
    "\n",
    "# Test \n",
    "test_mask = data[\"drug\", \"interacts_with\", \"gene\"].test_mask\n",
    "print(\"\\n===== Test Dataset =====\")\n",
    "print(\"Number of drug nodes:\", data['drug'].num_nodes)\n",
    "print(\"Number of gene nodes:\", data['gene'].num_nodes)\n",
    "print(\"Number of edges:\", test_mask.sum().item())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a3405fc0-42a5-4ca3-bc87-0b8452f9aee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.loader import DataLoader \n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class EdgeDataset(Dataset):\n",
    "    def __init__(self, data: HeteroData, mask: torch.Tensor):\n",
    "        self.edge_index = data[\"drug\", \"interacts_with\", \"gene\"].edge_index[:, mask]\n",
    "        self.edge_label = data[\"drug\", \"interacts_with\", \"gene\"].edge_label[mask]\n",
    "        self.num_edges = self.edge_index.size(1)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_edges\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.edge_index[:, idx], self.edge_label[idx]\n",
    "\n",
    "\n",
    "# Custom collate function to prevent automatic batching of edge_index\n",
    "def custom_collate_fn(batch):\n",
    "    edge_indices = [item[0] for item in batch]  # Extract edge_index from batch\n",
    "    edge_labels = [item[1] for item in batch]  # Extract edge_label from batch\n",
    "    \n",
    "    # Concatenate edge indices and labels\n",
    "    edge_index = torch.cat(edge_indices, dim=1)  # Shape: [2, batch_size]\n",
    "    edge_label = torch.cat(edge_labels, dim=0)   # Shape: [batch_size]\n",
    "\n",
    "    return edge_index, edge_label\n",
    "\n",
    "\n",
    "train_dataset = EdgeDataset(data, train_mask)\n",
    "valid_dataset = EdgeDataset(data, valid_mask)\n",
    "test_dataset = EdgeDataset(data, test_mask)\n",
    "\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, collate_fn=custom_collate_fn)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=128, shuffle=False, collate_fn=custom_collate_fn)\n",
    "test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, collate_fn=custom_collate_fn)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "887eda3b-b36b-4548-9229-48854ae99c2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of nodes:\n",
      "Drugs: 2726\n",
      "Genes: 2001\n",
      "$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$\n",
      "Train - Total: 4197, Positive: 1846, Negative: 2351\n",
      "Valid - Total: 600, Positive: 251, Negative: 349\n",
      "Test  - Total: 1200, Positive: 536, Negative: 664\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from collections import Counter\n",
    "\n",
    "# Function to count positive and negative edges\n",
    "def count_edges(edge_labels):\n",
    "    counter = Counter(edge_labels.tolist())\n",
    "    return counter[1.0], counter[0.0]  # Positive, Negative\n",
    "\n",
    "\n",
    "train_labels = train_dataset.edge_label\n",
    "valid_labels = valid_dataset.edge_label\n",
    "test_labels = test_dataset.edge_label\n",
    "\n",
    "\n",
    "print(\"Total number of nodes:\")\n",
    "print(f\"Drugs: {data['drug'].xs.size(0)}\")\n",
    "print(f\"Genes: {data['gene'].xs.size(0)}\")\n",
    "\n",
    "print(\"$\"*52)\n",
    "\n",
    "train_positive, train_negative = count_edges(train_labels)\n",
    "valid_positive, valid_negative = count_edges(valid_labels)\n",
    "test_positive, test_negative = count_edges(test_labels)\n",
    "\n",
    "\n",
    "\n",
    "print(f\"Train - Total: {train_labels.size(0)}, Positive: {train_positive}, Negative: {train_negative}\")\n",
    "print(f\"Valid - Total: {valid_labels.size(0)}, Positive: {valid_positive}, Negative: {valid_negative}\")\n",
    "print(f\"Test  - Total: {test_labels.size(0)}, Positive: {test_positive}, Negative: {test_negative}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f0ad2f69-4608-48b3-aa6d-25a6ebcc75ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drug LLM Input Size: 768\n",
      "Gene LLM Input Size: 1024\n",
      "Drug Structure Input Size: 1200\n",
      "Gene Structure Input Size: 1200\n"
     ]
    }
   ],
   "source": [
    "drug_l_input_size= data['drug']['xl'].size(1)\n",
    "gene_l_input_size= data['gene']['xl'].size(1)\n",
    "drug_s_input_size= data['drug']['xs'].size(1)\n",
    "gene_s_input_size= data['gene']['xs'].size(1)\n",
    "\n",
    "print(f\"Drug LLM Input Size: {drug_l_input_size}\")\n",
    "print(f\"Gene LLM Input Size: {gene_l_input_size}\")\n",
    "print(f\"Drug Structure Input Size: {drug_s_input_size}\")\n",
    "print(f\"Gene Structure Input Size: {gene_s_input_size}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a27d758f-be50-4562-aabe-5d35a572f9bd",
   "metadata": {},
   "source": [
    "# The below model takes structure and llm embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3717b68a-b44b-4540-9fcd-c60e5cbd16bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.nn import SAGEConv, to_hetero, BatchNorm\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import roc_auc_score, average_precision_score, accuracy_score, recall_score, confusion_matrix\n",
    "import os\n",
    "\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# it takes drug and gene llm and structure features and assigns weighs to them durung traning\n",
    "class FeatureFusion(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super().__init__()\n",
    "        self.fc = torch.nn.Linear(2 * hidden_channels, hidden_channels)\n",
    "\n",
    "    def forward(self, llm_features, struct_features):\n",
    "        concat_features = torch.cat([llm_features, struct_features], dim=-1)\n",
    "        alpha = torch.sigmoid(self.fc(concat_features)) # values are between 0 and 1\n",
    "        return alpha * llm_features + (1 - alpha) * struct_features # relative importance of each feature\n",
    "\n",
    "\n",
    "class Classifier(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super().__init__()\n",
    "        self.fc1 = torch.nn.Linear(2 * hidden_channels, hidden_channels)\n",
    "        self.fc2 = torch.nn.Linear(hidden_channels, 1)\n",
    "        self.dropout = torch.nn.Dropout(0.5)\n",
    "\n",
    "    def forward(self, drug_emb, gene_emb, edge_index):\n",
    "        if edge_index.shape[0] != 2:\n",
    "            edge_index = edge_index.t()\n",
    "\n",
    "        src, dst = edge_index\n",
    "        drug_emb_src = drug_emb[src]\n",
    "        gene_emb_dst = gene_emb[dst]\n",
    "\n",
    "        edge_emb = torch.cat([drug_emb_src, gene_emb_dst], dim=-1)\n",
    "        edge_emb = F.relu(self.fc1(edge_emb))\n",
    "        edge_emb = self.dropout(edge_emb)\n",
    "        return self.fc2(edge_emb).squeeze(-1)\n",
    "\n",
    "\n",
    "\n",
    "class GNN(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super().__init__()\n",
    "        self.conv1 = SAGEConv(hidden_channels, hidden_channels)\n",
    "        self.bn1 = BatchNorm(hidden_channels)\n",
    "        self.dropout = torch.nn.Dropout(0.5)\n",
    "        self.conv2 = SAGEConv(hidden_channels, hidden_channels)\n",
    "        self.bn2 = BatchNorm(hidden_channels)\n",
    "        \n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = F.relu(self.bn1(self.conv1(x, edge_index)))\n",
    "        x = self.dropout(x)\n",
    "        x = F.relu(self.bn2(self.conv2(x, edge_index)))\n",
    "        return x\n",
    "\n",
    "\n",
    "\n",
    "# Main Model\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels, drug_l_input_size, gene_l_input_size, drug_s_input_size, gene_s_input_size):\n",
    "        super().__init__()\n",
    "\n",
    "        # Drug \n",
    "        self.drug_llm_lin = torch.nn.Sequential(\n",
    "            torch.nn.Linear(drug_l_input_size, 512),\n",
    "            torch.nn.BatchNorm1d(512),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(512, hidden_channels),\n",
    "            torch.nn.BatchNorm1d(hidden_channels),\n",
    "            torch.nn.ReLU()\n",
    "        )\n",
    "        self.drug_struct_lin = torch.nn.Sequential(\n",
    "            torch.nn.Linear(drug_s_input_size, 1024),\n",
    "            torch.nn.BatchNorm1d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(1024, hidden_channels),\n",
    "            torch.nn.BatchNorm1d(hidden_channels),\n",
    "            torch.nn.ReLU()\n",
    "        )\n",
    "\n",
    "        # Gene \n",
    "        self.gene_llm_lin = torch.nn.Sequential(\n",
    "            torch.nn.Linear(gene_l_input_size, 1024),\n",
    "            torch.nn.BatchNorm1d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(1024, hidden_channels),\n",
    "            torch.nn.BatchNorm1d(hidden_channels),\n",
    "            torch.nn.ReLU()\n",
    "        )\n",
    "        self.gene_struct_lin = torch.nn.Sequential(\n",
    "            torch.nn.Linear(gene_s_input_size, 1024),\n",
    "            torch.nn.BatchNorm1d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(1024, hidden_channels),\n",
    "            torch.nn.BatchNorm1d(hidden_channels),\n",
    "            torch.nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.feature_fusion = FeatureFusion(hidden_channels)\n",
    "        self.gnn = GNN(hidden_channels)\n",
    "        self.gnn = to_hetero(self.gnn, metadata=data.metadata())\n",
    "        self.classifier = Classifier(hidden_channels)\n",
    "\n",
    "    def forward(self, data, edge_index, edge_label):\n",
    "        drug_xl = data[\"drug\"].xl.to(device).float()\n",
    "        drug_xs = data[\"drug\"].xs.to(device).float()\n",
    "        gene_xl = data[\"gene\"].xl.to(device).float()\n",
    "        gene_xs = data[\"gene\"].xs.to(device).float()\n",
    "\n",
    "        # Feature normalization\n",
    "        drug_xs = (drug_xs - drug_xs.mean(dim=0)) / (drug_xs.std(dim=0, unbiased=False) + 1e-8)\n",
    "        gene_xs = (gene_xs - gene_xs.mean(dim=0)) / (gene_xs.std(dim=0, unbiased=False) + 1e-8)\n",
    "\n",
    "        # Feature Fussion\n",
    "        drug_llm = self.drug_llm_lin(drug_xl)\n",
    "        drug_struct = self.drug_struct_lin(drug_xs)\n",
    "        drug_x = self.feature_fusion(drug_llm, drug_struct)\n",
    "\n",
    "        gene_llm = self.gene_llm_lin(gene_xl)\n",
    "        gene_struct = self.gene_struct_lin(gene_xs)\n",
    "        gene_x = self.feature_fusion(gene_llm, gene_struct)\n",
    "\n",
    "        # Pass through GNN\n",
    "        x_dict = {\"drug\": drug_x, \"gene\": gene_x}\n",
    "        x_dict = self.gnn(x_dict, data.edge_index_dict)\n",
    "\n",
    "        drug_emb = x_dict[\"drug\"]\n",
    "        gene_emb = x_dict[\"gene\"]\n",
    "        pred = self.classifier(drug_emb, gene_emb, edge_index)\n",
    "\n",
    "        return pred, edge_label\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "77de73ab-9f2c-4541-8c09-ad0858551f63",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run 1: AUROC = 0.9934, AUPRC = 0.9928, Accuracy = 0.9642, Sensitivity = 0.9813, Specificity = 0.9503\n",
      "Run 2: AUROC = 0.9929, AUPRC = 0.9915, Accuracy = 0.9617, Sensitivity = 0.9646, Specificity = 0.9593\n",
      "Run 3: AUROC = 0.9931, AUPRC = 0.9918, Accuracy = 0.9700, Sensitivity = 0.9646, Specificity = 0.9744\n",
      "Run 4: AUROC = 0.9921, AUPRC = 0.9911, Accuracy = 0.9708, Sensitivity = 0.9739, Specificity = 0.9684\n",
      "Run 5: AUROC = 0.9925, AUPRC = 0.9910, Accuracy = 0.9725, Sensitivity = 0.9701, Specificity = 0.9744\n",
      "\n",
      "Final Results:\n",
      "AUROC: Mean = 0.9928, Std = 0.0005\n",
      "AUPRC: Mean = 0.9916, Std = 0.0006\n",
      "Accuracy: Mean = 0.9678, Std = 0.0042\n",
      "Sensitivity: Mean = 0.9709, Std = 0.0063\n",
      "Specificity: Mean = 0.9654, Std = 0.0093\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import roc_auc_score, average_precision_score, accuracy_score, recall_score, confusion_matrix\n",
    "import numpy as np\n",
    "import gc\n",
    "import random\n",
    "\n",
    "# Training function\n",
    "\n",
    "def train_and_validate(model, data, train_loader, val_loader, device, patience=5, max_epochs=100, dataset_name=\"default\", run_id=1):\n",
    "    model_save_dir = f\"models/{dataset_name}\"\n",
    "    os.makedirs(model_save_dir, exist_ok=True)\n",
    "    \n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    data.edge_index_dict = {key: edge_index.to(device) for key, edge_index in data.edge_index_dict.items()}\n",
    "    model = model.to(device)\n",
    "\n",
    "    # Calculate positive weight for the loss function\n",
    "    all_labels = torch.cat([labels for _, labels in train_loader], dim=0)\n",
    "    negative_to_positive_ratio = len(all_labels[all_labels == 0]) / len(all_labels[all_labels == 1])  # Negatives/Positives\n",
    "    pos_weight=torch.tensor(negative_to_positive_ratio).to(device)\n",
    "\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005, weight_decay = 1e-2)\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max = 4, eta_min=1e-6)\n",
    "       \n",
    "    best_auroc = 0.0\n",
    "    best_val_loss = float(\"inf\")\n",
    "    epochs_no_improve = 0\n",
    "\n",
    "    for epoch in range(1, max_epochs + 1):\n",
    "        model.train()\n",
    "        total_train_loss, total_train_examples = 0, 0\n",
    "\n",
    "        for edge_index, edge_label in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            edge_index, edge_label = edge_index.to(device), edge_label.to(device)\n",
    "\n",
    "            pred, ground_truth = model(data, edge_index, edge_label)\n",
    "\n",
    "            loss = F.binary_cross_entropy_with_logits(pred, ground_truth, pos_weight)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            total_train_loss += loss.item() * pred.numel()\n",
    "            total_train_examples += pred.numel()\n",
    "\n",
    "        avg_train_loss = total_train_loss / total_train_examples\n",
    "\n",
    "        # Validation\n",
    "        model.eval()\n",
    "        total_val_loss, total_val_examples = 0, 0\n",
    "        preds, ground_truths = [], []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for edge_index, edge_label in val_loader:\n",
    "                edge_index, edge_label = edge_index.to(device), edge_label.to(device)\n",
    "                pred, ground_truth = model(data, edge_index, edge_label)\n",
    "                loss = F.binary_cross_entropy_with_logits(pred, ground_truth)\n",
    "\n",
    "                total_val_loss += loss.item() * pred.numel()\n",
    "                total_val_examples += pred.numel()\n",
    "\n",
    "                preds.append(pred.cpu())\n",
    "                ground_truths.append(ground_truth.cpu())\n",
    "\n",
    "        avg_val_loss = total_val_loss / total_val_examples\n",
    "        preds = torch.cat(preds, dim=0).sigmoid().numpy()\n",
    "        ground_truths = torch.cat(ground_truths, dim=0).numpy()\n",
    "\n",
    "        # Metrics\n",
    "        auroc = roc_auc_score(ground_truths, preds)\n",
    "\n",
    "        if auroc > best_auroc:\n",
    "            best_auroc = auroc\n",
    "            best_epoch = epoch\n",
    "            model_save_path = os.path.join(model_save_dir, f\"{dataset_name}_best_model_run_{run_number+1}.pt\")\n",
    "            torch.save(model.state_dict(), model_save_path)\n",
    "\n",
    "        if avg_val_loss < best_val_loss:\n",
    "            best_val_loss = avg_val_loss\n",
    "            epochs_no_improve = 0\n",
    "        else:\n",
    "            epochs_no_improve += 1\n",
    "\n",
    "        # Early stopping\n",
    "        if epochs_no_improve >= patience:\n",
    "            break\n",
    "\n",
    "    return best_auroc\n",
    "\n",
    " # Evaluate the model on the test set\n",
    "\n",
    "def evaluate_test_set(model, data, data_loader, device):\n",
    "    model.eval()\n",
    "    preds, ground_truths = [], []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for edge_index, edge_label in data_loader:\n",
    "            edge_index, edge_label = edge_index.to(device), edge_label.to(device)\n",
    "            pred, _ = model(data, edge_index, edge_label)\n",
    "            preds.append(pred.cpu())\n",
    "            ground_truths.append(edge_label.cpu())\n",
    "\n",
    "    preds = torch.cat(preds, dim=0).sigmoid().numpy()  \n",
    "    ground_truths = torch.cat(ground_truths, dim=0).numpy()\n",
    "\n",
    "    binary_preds = (preds > 0.6).astype(int)\n",
    "\n",
    "    auroc = roc_auc_score(ground_truths, preds)\n",
    "    aupr = average_precision_score(ground_truths, preds)\n",
    "    accuracy = accuracy_score(ground_truths, binary_preds)\n",
    "    sensitivity = recall_score(ground_truths, binary_preds)\n",
    "    tn, fp, fn, tp = confusion_matrix(ground_truths, binary_preds).ravel()\n",
    "    specificity = tn / (tn + fp)\n",
    "\n",
    "    return auroc, aupr, accuracy, sensitivity, specificity\n",
    "\n",
    "hidden_channels = 512\n",
    "\n",
    "model = Model(hidden_channels, drug_l_input_size, gene_l_input_size, drug_s_input_size, gene_s_input_size)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "auroc_scores, aupr_scores = [], []\n",
    "accuracy_scores, sensitivity_scores, specificity_scores = [], [], []\n",
    "\n",
    "for run_number in range(5):\n",
    "    # Clear memory and set random seed\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "    seed= random.randint(0, 2**32 - 1)  \n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "\n",
    "    #Train the model on the train and validation set\n",
    "    model = Model(hidden_channels, drug_l_input_size, gene_l_input_size, drug_s_input_size, gene_s_input_size)\n",
    "    train_and_validate(\n",
    "        model=model,\n",
    "        data=data,\n",
    "        train_loader=train_loader,\n",
    "        val_loader=valid_loader,\n",
    "        device=device,\n",
    "        dataset_name=sub_dataset,\n",
    "        run_id=run_number\n",
    "    )\n",
    "\n",
    "   \n",
    "    model_path = f\"models/{sub_dataset}/{sub_dataset}_best_model_run_{run_number+1}.pt\"\n",
    "    model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "    auroc, aupr, accuracy, sensitivity, specificity = evaluate_test_set(model, data, test_loader, device)\n",
    "\n",
    "    auroc_scores.append(auroc)\n",
    "    aupr_scores.append(aupr)\n",
    "    accuracy_scores.append(accuracy)\n",
    "    sensitivity_scores.append(sensitivity)\n",
    "    specificity_scores.append(specificity)\n",
    "\n",
    "    print(f\"Run {run_number+1}: AUROC = {auroc:.4f}, AUPRC = {aupr:.4f}, Accuracy = {accuracy:.4f}, Sensitivity = {sensitivity:.4f}, Specificity = {specificity:.4f}\")\n",
    "\n",
    "# Calculate mean and standard deviation\n",
    "print(\"\\nFinal Results:\")\n",
    "print(f\"AUROC: Mean = {np.mean(auroc_scores):.4f}, Std = {np.std(auroc_scores):.4f}\")\n",
    "print(f\"AUPRC: Mean = {np.mean(aupr_scores):.4f}, Std = {np.std(aupr_scores):.4f}\")\n",
    "print(f\"Accuracy: Mean = {np.mean(accuracy_scores):.4f}, Std = {np.std(accuracy_scores):.4f}\")\n",
    "print(f\"Sensitivity: Mean = {np.mean(sensitivity_scores):.4f}, Std = {np.std(sensitivity_scores):.4f}\")\n",
    "print(f\"Specificity: Mean = {np.mean(specificity_scores):.4f}, Std = {np.std(specificity_scores):.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d69ca68b-05d8-49df-a613-67a9d69f2d87",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}