biobridge_primekg.py

"""
Class for loading BioBridgePrimeKG dataset.
"""

import os
import pickle
import json
import requests
import numpy as np
import pandas as pd
from tqdm import tqdm
from .dataset import Dataset
from .primekg import PrimeKG

class BioBridgePrimeKG(Dataset):
    """
    Class for loading BioBridgePrimeKG dataset.
    It downloads the data from the BioBridge repo and stores it in the local directory.
    The data is then loaded into pandas DataFrame of nodes and edges.
    This class was adapted from the BioBridge repo:
    https://github.com/RyanWangZf/BioBridge
    """

    def __init__(self,
                 primekg_dir: str = "../../../data/primekg/",
                 local_dir: str = "../../../data/biobridge_primekg/",
                 random_seed: int=0,
                 n_neg_samples: int=5):
        """
        Constructor for BioBridgePrimeKG class.

        Args:
            primekg_dir (str): The directory of PrimeKG dataset.
            local_dir (str): The directory to store the downloaded data.
            random_seed (int): The random seed value.
        """
        self.name: str = "biobridge_primekg"
        self.primekg_dir: str = primekg_dir
        self.local_dir: str = local_dir
        self.random_seed = random_seed
        self.n_neg_samples = n_neg_samples
        # Preselected node types:
        # protein, molecular function, cellular component, biological process, drug, disease
        self.preselected_node_types = ["protein", "mf", "cc", "bp", "drug", "disease"]
        self.node_type_map = {
            "protein": "gene/protein",
            "mf": "molecular_function",
            "cc": "cellular_component",
            "bp": "biological_process",
            "drug": "drug",
            "disease": "disease",
        }

        # Attributes to store the data
        self.primekg = None
        self.primekg_triplets = None
        self.primekg_triplets_negative = None
        self.data_config = None
        self.emb_dict = None
        self.df_train = None
        self.df_node_train = None
        self.df_test = None
        self.df_node_test = None
        self.node_info_dict = None

        # Set up the dataset
        self.setup()

    def setup(self):
        """
        A method to set up the dataset.
        """
        # Make the directories if it doesn't exist
        os.makedirs(os.path.dirname(self.primekg_dir), exist_ok=True)
        os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)

        # Set the random seed
        self.set_random_seed(self.random_seed)

        # Set SettingWithCopyWarning  warnings to none
        pd.options.mode.chained_assignment = None

    def _load_primekg(self) -> PrimeKG:
        """
        Private method to load related files of PrimeKG dataset.

        Returns:
            PrimeKG: The PrimeKG dataset.
        """
        primekg_data = PrimeKG(local_dir=self.primekg_dir)
        primekg_data.load_data()

        return primekg_data

    def _download_file(self,
                       remote_url:str,
                       local_dir: str,
                       local_filename: str):
        """
        A helper function to download a file from remote URL to the local directory.

        Args:
            remote_url (str): The remote URL of the file to be downloaded.
            local_dir (str): The local directory to store the downloaded file.
            local_filename (str): The local filename to store the downloaded file.
        """
        # Make the local directory if it does not exist
        if not os.path.exists(local_dir):
            os.makedirs(local_dir)
        # Download the file from remote URL to local directory
        local_path = os.path.join(local_dir, local_filename)
        if os.path.exists(local_path):
            print(f"File {local_filename} already exists in {local_dir}.")
        else:
            print(f"Downloading {local_filename} from {remote_url} to {local_dir}...")
            response = requests.get(remote_url, stream=True, timeout=300)
            response.raise_for_status()
            progress_bar = tqdm(
                total=int(response.headers.get("content-length", 0)),
                unit="iB",
                unit_scale=True,
            )
            with open(os.path.join(local_dir, local_filename), "wb") as file:
                for data in response.iter_content(1024):
                    progress_bar.update(len(data))
                    file.write(data)
            progress_bar.close()

    def _load_data_config(self) -> dict:
        """
        Load the data config file of BioBridgePrimeKG dataset.

        Returns:
            dict: The data config file of BioBridgePrimeKG dataset.
        """
        # Download the data config file of BioBridgePrimeKG
        self._download_file(
            remote_url= ('https://raw.githubusercontent.com/RyanWangZf/BioBridge/'
                         'refs/heads/main/data/BindData/data_config.json'),
            local_dir=self.local_dir,
            local_filename='data_config.json')

        # Load the downloaded data config file
        with open(os.path.join(self.local_dir, 'data_config.json'), 'r', encoding='utf-8') as f:
            data_config = json.load(f)

        return data_config

    def _build_node_embeddings(self) -> dict:
        """
        Build the node embeddings for BioBridgePrimeKG dataset.

        Returns:
            dict: The dictionary of node embeddings.
        """
        processed_file_path = os.path.join(self.local_dir, "embeddings", "embedding_dict.pkl")
        if os.path.exists(processed_file_path):
            # Load the embeddings from the local directory
            with open(processed_file_path, "rb") as f:
                emb_dict_all = pickle.load(f)
        else:
            # Download the embeddings from the BioBridge repo and further process them
            # List of embedding source files
            url = ('https://media.githubusercontent.com/media/RyanWangZf/BioBridge/'
                   'refs/heads/main/data/embeddings/esm2b_unimo_pubmedbert/')
            file_list = [f"{n}.pkl" for n in self.preselected_node_types]

            # Download the embeddings
            for file in file_list:
                self._download_file(remote_url=os.path.join(url, file),
                                    local_dir=os.path.join(self.local_dir, "embeddings"),
                                    local_filename=file)

            # Unified embeddings
            emb_dict_all = {}
            for file in file_list:
                with open(os.path.join(self.local_dir, "embeddings", file), "rb") as f:
                    emb = pickle.load(f)
                emb_ar = emb["embedding"]
                if not isinstance(emb_ar, list):
                    emb_ar = emb_ar.tolist()
                emb_dict_all.update(dict(zip(emb["node_index"], emb_ar)))

            # Store embeddings
            with open(processed_file_path, "wb") as f:
                pickle.dump(emb_dict_all, f)

        return emb_dict_all

    def _build_full_triplets(self) -> pd.DataFrame:
        """
        Build the full triplets for BioBridgePrimeKG dataset.

        Returns:
            pd.DataFrame: The full triplets for BioBridgePrimeKG dataset.
        """
        processed_file_path = os.path.join(self.local_dir, "processed", "triplet_full.tsv.gz")
        if os.path.exists(processed_file_path):
            # Load the file from the local directory
            with open(processed_file_path, "rb") as f:
                primekg_triplets = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)

            # Load each dataframe in the local directory
            node_info_dict = {}
            for i, node_type in enumerate(self.preselected_node_types):
                with open(os.path.join(self.local_dir, "processed",
                                       f"{node_type}.csv"), "rb") as f:
                    df_node = pd.read_csv(f)
                node_info_dict[self.node_type_map[node_type]] = df_node
        else:
            # Download the related files from the BioBridge repo and further process them
            # List of processed files
            url = ('https://media.githubusercontent.com/media/RyanWangZf/BioBridge/'
                   'refs/heads/main/data/Processed/')
            file_list = ["protein", "molecular", "cellular", "biological", "drug", "disease"]

            # Download the processed files
            for i, file in enumerate(file_list):
                self._download_file(remote_url=os.path.join(url, f"{file}.csv"),
                                    local_dir=os.path.join(self.local_dir, "processed"),
                                    local_filename=f"{self.preselected_node_types[i]}.csv")

            # Build the node index list
            node_info_dict = {}
            node_index_list = []
            for i, file in enumerate(file_list):
                df_node = pd.read_csv(os.path.join(self.local_dir, "processed",
                                                   f"{self.preselected_node_types[i]}.csv"))
                node_info_dict[self.node_type_map[self.preselected_node_types[i]]] = df_node
                node_index_list.extend(df_node["node_index"].tolist())

            # Filter the PrimeKG dataset to take into account only the selected node types
            primekg_triplets = self.primekg.get_edges().copy()
            primekg_triplets = primekg_triplets[
                primekg_triplets["head_index"].isin(node_index_list) &\
                primekg_triplets["tail_index"].isin(node_index_list)
            ]
            primekg_triplets = primekg_triplets.reset_index(drop=True)

            # Perform mapping of node types
            primekg_triplets["head_type"] = primekg_triplets["head_type"].apply(
                lambda x: self.data_config["node_type"][x]
            )
            primekg_triplets["tail_type"] = primekg_triplets["tail_type"].apply(
                lambda x: self.data_config["node_type"][x]
            )

            # Perform mapping of relation types
            primekg_triplets["display_relation"] = primekg_triplets["display_relation"].apply(
                lambda x: self.data_config["relation_type"][x]
            )

            # Store the processed triplets
            primekg_triplets.to_csv(processed_file_path, sep="\t", compression="gzip", index=False)

        return primekg_triplets, node_info_dict

    def _build_train_test_split(self):
        """
        Build the train-test split for BioBridgePrimeKG dataset.
        """
        if os.path.exists(os.path.join(self.local_dir, "processed",
                                       "triplet_full_altered.tsv.gz")):
            # Load each dataframe in the local directory
            with open(os.path.join(self.local_dir, "processed",
                                   "triplet_train.tsv.gz"), "rb") as f:
                df_train = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)

            with open(os.path.join(self.local_dir, "processed",
                                   "node_train.tsv.gz"), "rb") as f:
                df_node_train = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)

            with open(os.path.join(self.local_dir, "processed",
                                   "triplet_test.tsv.gz"), "rb") as f:
                df_test = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)

            with open(os.path.join(self.local_dir, "processed",
                                   "node_test.tsv.gz"), "rb") as f:
                df_node_test = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)

            with open(os.path.join(self.local_dir, "processed",
                                   "triplet_full_altered.tsv.gz"), "rb") as f:
                triplets = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
        else:
            # Filtering out some nodes in the embedding dictionary
            triplets = self.primekg_triplets.copy()
            triplets = triplets[
                triplets["head_index"].isin(list(self.emb_dict.keys())) &\
                triplets["tail_index"].isin(list(self.emb_dict.keys()))
            ].reset_index(drop=True)

            # Perform splitting of the triplets
            list_split = {
                "train": [],
                "test": [],
            }
            node_split = {
                "train": {
                    "node_index": [],
                    "node_type": [],
                },
                "test": {
                    "node_index": [],
                    "node_type": [],
                }
            }
            # Loop over the node types
            for node_type in triplets["head_type"].unique():
                df_sub = triplets[triplets["head_type"] == node_type]
                all_x_indexes = df_sub["head_index"].unique()
                # By default, we use 90% of the nodes for training and 10% for testing
                te_x_indexes = np.random.choice(
                    all_x_indexes, size=int(0.1*len(all_x_indexes)), replace=False
                )
                df_subs = {}
                df_subs["test"] = df_sub[df_sub["head_index"].isin(te_x_indexes)]
                df_subs["train"] = df_sub[~df_sub["head_index"].isin(te_x_indexes)]
                list_split["train"].append(df_subs["train"])
                list_split["test"].append(df_subs["test"])

                # record the split
                node_index = {}
                node_index["train"] = df_subs["train"]["head_index"].unique()
                node_split["train"]["node_index"].extend(node_index["train"].tolist())
                node_split["train"]["node_type"].extend([node_type]*len(node_index["train"]))
                node_index["test"] = df_subs["test"]["head_index"].unique()
                node_split["test"]["node_index"].extend(node_index["test"].tolist())
                node_split["test"]["node_type"].extend([node_type]*len(node_index["test"]))

                print(f"Number of {node_type} nodes in train: {len(node_index["train"])}")
                print(f"Number of {node_type} nodes in test: {len(node_index["test"])}")

            # Prepare train and test DataFrames
            df_train = pd.concat(list_split["train"])
            df_node_train = pd.DataFrame(node_split["train"])
            df_test = pd.concat(list_split["test"])
            df_node_test = pd.DataFrame(node_split["test"])

            # Store each dataframe in the local directory
            df_train.to_csv(os.path.join(self.local_dir, "processed", "triplet_train.tsv.gz"),
                            sep="\t", compression="gzip", index=False)
            df_node_train.to_csv(os.path.join(self.local_dir, "processed", "node_train.tsv.gz"),
                                sep="\t", compression="gzip", index=False)
            df_test.to_csv(os.path.join(self.local_dir, "processed", "triplet_test.tsv.gz"),
                           sep="\t", compression="gzip", index=False)
            df_node_test.to_csv(os.path.join(self.local_dir, "processed", "node_test.tsv.gz"),
                                sep="\t", compression="gzip", index=False)
            # Store altered full triplets as well
            triplets.to_csv(os.path.join(self.local_dir, "processed",
                                         "triplet_full_altered.tsv.gz"),
                            sep="\t", compression="gzip", index=False)

        return df_train, df_node_train, df_test, df_node_test, triplets

    # def _negative_sampling(self,
    #                        batch_df: pd.DataFrame,
    #                        process_index: int,
    #                        index_map: dict,
    #                        node_train_dict: dict) -> pd.DataFrame:
    #     """
    #     A helper function to perform negative sampling for a batch of triplets.
    #     """
    #     negative_y_index_list = []
    #     for _, row in tqdm(batch_df.iterrows(),
    #                        total=batch_df.shape[0],
    #                        desc=f"Process {process_index}"):
    #         x_index = row['head_index']
    #         # y_index = row['y_index']
    #         y_index_type = row['tail_type']
    #         paired_y_index_list = index_map[x_index]

    #         # sample a list of negative y_index
    #         node_train_sub = node_train_dict[y_index_type]
    #         negative_y_index = node_train_sub[
    #             ~node_train_sub['node_index'].isin(paired_y_index_list)
    #         ]['node_index'].sample(self.n_neg_samples).tolist()
    #         negative_y_index_list.append(negative_y_index)

    #     batch_df.loc[:, 'negative_tail_index'] = negative_y_index_list
    #     return batch_df

    # def _build_negative_triplets(self,
    #                              chunk_size: int=100000,
    #                              n_neg_samples: int=10):
    #     """
    #     Build the negative triplets for BioBridgePrimeKG dataset.
    #     """
    #     processed_file_path = os.path.join(self.local_dir,
    #                                        "processed",
    #                                        "triplet_train_negative.tsv.gz")
    #     if os.path.exists(processed_file_path):
    #         # Load the negative triplets from the local directory
    #         with open(processed_file_path, "rb") as f:
    #             triplets_negative = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
    #     else:
    #         # Set the number samples for negative sampling
    #         self.n_neg_samples = n_neg_samples

    #         # Split node list by type
    #         node_train_dict = {}
    #         type_list = self.df_node_train['node_type'].unique()
    #         for node_type in type_list:
    #             node_train_dict[node_type] = self.df_node_train[
    #                 self.df_node_train['node_type'] == node_type
    #             ].reset_index(drop=True)

    #         # create an index mapping from x_index to y_index
    #         index_map = self.df_train[
    #             ['head_index', 'tail_index']
    #         ].drop_duplicates().groupby('head_index').agg(list).to_dict()['tail_index']

    #         # Negative sampling
    #         batch_df_list = []
    #         for i in tqdm(range(0, self.df_train.shape[0], chunk_size)):
    #             batch_df_list.append(self.df_train.iloc[i:i+chunk_size])
    #         # Process negative sampling
    #         results = [
    #             self._negative_sampling(batch_df,
    #                                     num_piece,
    #                                     index_map,
    #                                     node_train_dict)
    #                                     for num_piece, batch_df in enumerate(batch_df_list)
    #         ]

    #         # Store the negative triplets
    #         triplets_negative = pd.concat(results, axis=0)
    #         triplets_negative.to_csv(processed_file_path,
    #                                  sep="\t", compression="gzip", index=False)

    #     # Set attribute
    #     self.primekg_triplets_negative = triplets_negative

    #     return triplets_negative

    # def load_data(self,
    #               build_neg_triplest: bool= False,
    #               chunk_size: int=100000,
    #               n_neg_samples: int=10):

    def load_data(self):
        """
        Load the BioBridgePrimeKG dataset into pandas DataFrame of nodes and edges.

        Args:
            build_neg_triplest (bool): Whether to build negative triplets.
            chunk_size (int): The chunk size for negative sampling.
            n_neg_samples (int): The number of negative samples for negative sampling.
        """
        # Load PrimeKG dataset
        print("Loading PrimeKG dataset...")
        self.primekg = self._load_primekg()

        # Load data config file of BioBridgePrimeKG
        print("Loading data config file of BioBridgePrimeKG...")
        self.data_config = self._load_data_config()

        # Build node embeddings
        print("Building node embeddings...")
        self.emb_dict = self._build_node_embeddings()

        # Build full triplets
        print("Building full triplets...")
        self.primekg_triplets, self.node_info_dict = self._build_full_triplets()

        # Build train-test split
        print("Building train-test split...")
        self.df_train, self.df_node_train, self.df_test, self.df_node_test, self.primekg_triplets =\
        self._build_train_test_split()

        # if build_neg_triplest:
        #     # Build negative triplets
        #     print("Building negative triplets...")
        #     self.primekg_triplets_negative = self._build_negative_triplets(
        #         chunk_size=chunk_size,
        #         n_neg_samples=n_neg_samples
        #     )

    def set_random_seed(self, seed: int):
        """
        Set the random seed for reproducibility.

        Args:
            seed (int): The random seed value.
        """
        np.random.seed(seed)

    def get_primekg(self) -> PrimeKG:
        """
        Get the PrimeKG dataset.

        Returns:
            PrimeKG: The PrimeKG dataset.
        """
        return self.primekg

    def get_data_config(self) -> dict:
        """
        Get the data config file of BioBridgePrimeKG dataset.

        Returns:
            dict: The data config file of BioBridgePrimeKG dataset.
        """
        return self.data_config

    def get_node_embeddings(self) -> dict:
        """
        Get the node embeddings for BioBridgePrimeKG dataset.

        Returns:
            dict: The dictionary of node embeddings.
        """
        return self.emb_dict

    def get_primekg_triplets(self) -> pd.DataFrame:
        """
        Get the full triplets for BioBridgePrimeKG dataset.

        Returns:
            pd.DataFrame: The full triplets for BioBridgePrimeKG dataset.
        """
        return self.primekg_triplets

    # def get_primekg_triplets_negative(self) -> pd.DataFrame:
    #     """
    #     Get the negative triplets for BioBridgePrimeKG dataset.

    #     Returns:
    #         pd.DataFrame: The negative triplets for BioBridgePrimeKG dataset.
    #     """
    #     return self.primekg_triplets_negative

    def get_train_test_split(self) -> dict:
        """
        Get the train-test split for BioBridgePrimeKG dataset.

        Returns:
            dict: The train-test split for BioBridgePrimeKG dataset.
        """
        return {
            "train": self.df_train,
            "node_train": self.df_node_train,
            "test": self.df_test,
            "node_test": self.df_node_test
        }

    def get_node_info_dict(self) -> dict:
        """
        Get the node information dictionary for BioBridgePrimeKG dataset.

        Returns:
            dict: The node information dictionary for BioBridgePrimeKG dataset.
        """
        return self.node_info_dict