Oncodrive3D 1.0.4__py3-none-any.whl

scripts/datasets/utils.py

@@ -0,0 +1,394 @@
+"""
+Module including a collection of functions that provide
+general-purpose functionalities that can be used across
+different parts of the dataset building process.
+"""
+
+import gzip
+import hashlib
+import io
+import os
+import tarfile
+import time
+from zipfile import ZipFile
+
+import logging
+import daiquiri
+import numpy as np
+import pandas as pd
+import requests
+from Bio import SeqIO
+from pypdl import Pypdl as Downloader
+import aiohttp
+# from unipressed import IdMappingClient
+
+from scripts import __logger_name__
+
+logger = daiquiri.getLogger(__logger_name__ + ".build.utils")
+
+
+# General utils
+
+def rounded_up_division(num, den):
+    """
+    Perform a division and round up the result.
+    """
+
+    return -(-num // den)
+
+
+def get_pos_fragments(mut_gene_df):
+    """
+    Get the corresponding fragment of each position of the protein.
+    """
+
+    max_f = rounded_up_division(max(mut_gene_df["Pos"]), 1400)
+    bins = [n * 1400 for n in range(0, max_f+1)]
+    group_names = list(range(1, max_f+1))
+
+    return pd.cut(mut_gene_df["Pos"], bins, labels = group_names)
+
+
+def get_species(species):
+    """
+    Change species name to accepted format.
+    """
+
+    if species == "human" or species.capitalize() == "Homo sapiens":
+        species = "Homo sapiens"
+    elif species == "mouse" or species.capitalize() == "Mus musculus":
+        species = "Mus musculus"
+    else:
+        raise RuntimeError(
+            f"Failed to recognize '{species}' as species. Currently accepted ones are 'Homo sapiens' and 'Mus musculus'. Exiting.."
+            )
+
+    return species
+
+
+# Download
+
+CHECKSUM = {
+    "UP000005640_9606_HUMAN_v4": "bf62d5402cb1c4580d219335a9af1ac831416edfbf2892391c8197a8356091f2",
+    "UP000000589_10090_MOUSE_v4" : "eb6c529c8757d511b75f4856c1a789378478e6255a442517ad8579708787bbab",
+    "mane_overlap_v4" : "c01e9b858c5415cfe2eae7e52a561aa8a872ba0d5d4891ba0cec327b3af49d69"
+}
+
+
+def assert_proteome_integrity(file_path, proteome):
+
+    if proteome in CHECKSUM.keys():
+        logger.debug('Asserting integrity of file..')
+        try:
+            if CHECKSUM[proteome] == calculate_hash(file_path):
+                logger.debug('File integrity check: PASS')
+                return "PASS"
+            else:
+                logger.debug('File integrity check: FAIL')
+                return "FAIL"
+        except Exception as e:
+            logger.debug('File integrity check: FAIL')
+            logger.debug(f'Error: {e}')
+            return "FAIL"
+    else:
+        logger.warning("Assertion skipped: Proteome checksum not in records.")
+        return "PASS"
+
+
+def calculate_hash(filepath: str, hash_func=hashlib.sha256) -> str:
+    """
+    Calculate the hash of a file using the specified hash function.
+
+    Args:
+        filepath (str): The path to the file for which to calculate the hash.
+        hash_func (hashlib.Hash): The hash function to use (default is hashlib.sha256).
+
+    Returns:
+        str: The hexadecimal representation of the calculated hash.
+    """
+
+    with open(filepath, 'rb') as file:
+        hash_obj = hash_func()
+        for chunk in iter(lambda: file.read(8192), b''):
+            hash_obj.update(chunk)
+    return hash_obj.hexdigest()
+
+
+def download_single_file(url: str, destination: str, threads: int, proteome=None) -> None:
+    """
+    Downloads a file from a URL and saves it to the specified destination.
+
+    Args:
+        url (str): The URL of the file to download.
+        destination (str): The local path where the file will be saved.
+    """
+
+    num_connections = 15 if threads > 40 else threads
+
+    if os.path.exists(destination):
+        logger.debug(f"File {destination} already exists..")
+        if proteome is not None:
+            status = assert_proteome_integrity(destination, proteome)
+            if status == "PASS":
+                logger.debug(f"File {destination} already exists: Skipping download..")
+                return None
+            else:
+                logger.debug(f"File {destination} already exists but failed integrity check: Retrying download..")
+
+    logger.debug(f'Downloading {url}')
+    logger.debug(f"Downloading to {destination}")
+    dl = Downloader(timeout=aiohttp.ClientTimeout(sock_read=400), ssl=False)
+    dl.start(url, destination, segments=num_connections, display=True, retries=10, clear_terminal=False)
+    logger.debug('Download complete')
+
+
+def extract_tar_file(file_path, path):
+
+    checkpoint = os.path.join(path, ".tar_checkpoint.txt")
+
+    if os.path.exists(checkpoint):
+         logger.debug('Tar already extracted: Skipping..')
+    else:
+        with tarfile.open(file_path, "r") as tar:
+            tar.extractall(path)
+            logger.debug(f'Extracted { int(len(tar.getnames())/2)} structure.')
+            with open(checkpoint, "w") as f:
+                f.write('')
+
+
+def extract_zip_file(file_path, path):
+    checkpoint = os.path.join(path, ".zip_checkpoint.txt")
+
+    if os.path.exists(checkpoint):
+        logger.debug('ZIP file already extracted: Skipping..')
+    else:
+        with ZipFile(file_path, "r") as zip_file:
+            zip_file.extractall(path)
+            logger.debug(f'Extracted {len(zip_file.namelist())} files.')
+            with open(checkpoint, "w") as f:
+                f.write('')
+
+
+# PDB
+
+def get_af_id_from_pdb(path_structure):
+    """
+    Get AlphaFold 2 identifier (UniprotID_F) from path
+    """
+
+    return path_structure.split("AF-")[1].split("-model")[0]
+
+
+def get_seq_from_pdb(path_structure):
+    """
+    Get sequense of amino acid residues from PDB structure.
+    """
+    if path_structure.endswith("gz"):
+        with gzip.open(path_structure, 'rt') as handle:
+            return np.array([record.seq for record in SeqIO.parse(handle, 'pdb-seqres')][0])
+    else:
+        with open(path_structure, 'r') as handle:
+            return np.array([record.seq for record in SeqIO.parse(handle, 'pdb-seqres')][0])
+
+
+def get_pdb_path_list_from_dir(path_dir):
+    """
+    Takes as input a path of a given directory and it
+    outputs a list of paths of the contained PDB files.
+    """
+
+    pdb_files = os.listdir(path_dir)
+    pdb_path_list = [f"{os.path.join(path_dir, f)}" for f in pdb_files if ".pdb" in f and not os.path.isdir(os.path.join(path_dir, f))]
+    return pdb_path_list
+
+
+# Uniprot ID to Hugo symbols mapping
+
+def split_lst_into_chunks(lst, batch_size = 5000):
+    """
+    Simple split a list into list of list of chunk_size elements.
+    """
+
+    return [lst[i:i+batch_size] for i in range(0, len(lst), batch_size)]
+
+
+def convert_dict_hugo_to_uniprot(dict_uniprot_hugo):
+    """
+    Convert a Uniprot IDs to Hugo symbol dictionary to a Hugo symbo to
+    Uniprot IDs dictionary, if multiple Hugo symbols are mapped to the
+    same Uniprot ID, add them as multiple keys.
+    """
+
+    dict_hugo_uniprot = {}
+
+    for uni_id, gene in dict_uniprot_hugo.items():
+        if type(gene) == str:
+            for g in gene.split(" "):
+                dict_hugo_uniprot[g] = uni_id
+
+    return dict_hugo_uniprot
+
+
+def uniprot_to_hudo_df(uniprot_ids):
+    """
+    Given a list of Uniprot IDs (from any species), request an Id
+    mapping job to UniprotKB to retrieve the corresponding Hugo
+    symbols and additional protein info. Return a pandas dataframe.
+    It is recommended to provide batches of IDs up to 5000 elements.
+    """
+
+    job_id = get_mapping_jobid(uniprot_ids)
+    url = f"https://rest.uniprot.org/idmapping/uniprotkb/results/stream/{job_id}?compressed=true&fields=accession%2Cid%2Cprotein_name%2Cgene_names%2Corganism_name%2Clength&format=tsv"
+    df = load_df_from_url(url)
+
+    i = 60
+    while df is None:
+        time.sleep(1)
+        df = load_df_from_url(url)
+        if i % 180 == 0:
+            logger.debug("Waiting for UniprotKB mapping job to produce url..")
+        i += 1
+
+    return df
+
+
+def load_df_from_url(url):
+    """
+    Load a pandas dataframe from url.
+    """
+
+    try:
+        response = requests.get(url, timeout=(160, 160))
+        decompressed_data = gzip.decompress(response.content)
+        df = pd.read_csv(io.BytesIO(decompressed_data), sep='\t')
+    except:
+        df = None
+
+    return df
+
+
+def get_response_jobid(response):
+    """
+    Get jobId after submitting ID Mapping job to UniprotKB.
+    """
+
+    try:
+        data = response.json()
+        job_id = data.get("jobId")
+    except:
+        job_id = None
+
+    return job_id
+
+
+def get_mapping_jobid(uniprot_ids):
+    """
+    Submit an ID Mapping job to UniprotKB.
+    """
+
+    command = f"https://rest.uniprot.org/idmapping/run?from=UniProtKB_AC-ID&to=UniProtKB&ids={','.join(uniprot_ids)}"
+
+    response = "INIT"
+    while str(response) != "<Response [200]>":
+        if response != "INIT":
+            time.sleep(10)
+        try:
+            response = requests.post(command)
+        except requests.exceptions.RequestException as e:
+            response = "ERROR"
+            logger.debug(f"Request failed {e}: Retrying")
+
+    job_id = get_response_jobid(response)
+    i = 60
+    while job_id is None:
+        time.sleep(1)
+        job_id = get_response_jobid(response)
+        if i % 60 == 0:
+            logger.debug("Requesting ID mapping job to UniprotKB for IDs..")
+        i += 1
+
+    return job_id
+
+
+def uniprot_to_hugo(uniprot_ids, hugo_as_keys=False, batch_size=5000):
+    """
+    Given a list of Uniprot IDs (any species.), request an Id mapping
+    job to UniprotKB to retrieve the corresponding Hugo symbols.
+    Return a dictionary of Uniprot IDs to Hugo symbols or vice versa.
+    """
+
+    # Split uniprot IDs into chunks
+    uniprot_ids_lst = split_lst_into_chunks(uniprot_ids, batch_size)
+
+    # Get a dataframe including all IDs mapping info
+    df_lst = []
+    for i, ids in enumerate(uniprot_ids_lst):
+        logger.debug(f"Batch {i+1}/{len(uniprot_ids_lst)} ({len(ids)} IDs)..")
+        df = uniprot_to_hudo_df(ids)
+        df_lst.append(df)
+    df = pd.concat(df_lst)
+
+    # Get a dictionary for Uniprot ID to Hugo symbols
+    dictio = {}
+    for i, r in df[["Entry", "Gene Names"]].iterrows():
+        uni_id, gene = r
+        dictio[uni_id] = gene
+
+    # Convert to a dictionary of Hugo symbols to Uniprot IDs
+    if hugo_as_keys:
+        dictio = convert_dict_hugo_to_uniprot(dictio)
+
+    return dictio
+
+
+# def uniprot_to_hugo_pressed(uniprot_ids, hugo_as_keys=False, batch_size=5000, max_attempts=15):
+#     """
+#     Given a list of Uniprot IDs (any species.), return a
+#     dictionary of Uniprot IDs to Hugo symbols or vice versa.
+#     """
+
+#     # Split uniprot IDs into chunks
+#     uniprot_ids_lst = split_lst_into_chunks(uniprot_ids, batch_size)
+
+#     # Get a dataframe including all IDs mapping info
+#     result_lst = []
+
+#     for i, ids in enumerate(uniprot_ids_lst):
+#         logger.debug(f"Batch {i+1}/{len(uniprot_ids_lst)} ({len(ids)} IDs)..")
+#         status = "INIT"
+
+#         n = 0
+#         while status != "FINISHED":
+#             try:
+#                 request = IdMappingClient.submit(source="UniProtKB_AC-ID",
+#                                                  dest="Gene_Name",
+#                                                  ids={uni_id for uni_id in ids})
+#                 j = 0
+#                 while status != "FINISHED":
+#                     time.sleep(15)
+#                     status = request.get_status()
+#                     if status == "FINISHED":
+#                         result_lst.append(list(request.each_result()))
+#                     else:
+#                         logger.debug(f"Waiting for UniprotKB to process the job..")
+#                         j += 1
+#                         if j == max_attempts:
+#                             logger.debug(f"Failed to obtain Uniprot ID to Hugo Symbol mapping: Retrying..")
+#                             sys.exit()
+#             except Exception as e:
+#                 n += 1
+#                 if n == max_attempts:
+#                     logger.error(f"Failed to obtain Uniprot ID to Hugo Symbol mapping and reached maximum attempts: Exiting..")
+#                     sys.exit()
+#                 else:
+#                     status = "ERROR"
+#                     logger.debug(f"Error while obtaining Uniprot ID to Hugo Symbol mapping {e}: Retrying..")
+
+#     result_lst = [entry for batch in result_lst for entry in batch]
+#     result_dict = {r["from"] : (r["to"].split()[0] if len(r["to"].split()) > 1 else r["to"]) for r in result_lst}
+
+#     # Convert to a dictionary of Hugo symbols to Uniprot IDs
+#     if hugo_as_keys:
+#         result_dict = convert_dict_hugo_to_uniprot(result_dict)
+
+#     return result_dict

scripts/globals.py

@@ -0,0 +1,169 @@
+import os
+import logging
+import daiquiri
+import click
+import shutil
+from datetime import datetime
+
+from functools import wraps
+
+from scripts import __logger_name__
+
+
+logger = daiquiri.getLogger(__logger_name__)
+
+DATE = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+FORMAT = "%(asctime)s - %(color)s%(levelname)-7s%(color_stop)s | %(name)s - %(color)s%(message)s%(color_stop)s"
+
+
+# =========
+#  Logging
+# =========
+
+def setup_logging_decorator(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        log_dir = os.path.join(click.get_current_context().params['output_dir'], 'log')
+        command_name = click.get_current_context().command.name
+
+        if command_name in ['run', 'plot', 'chimerax_plot', 'chimerax-plot']:
+            cohort = click.get_current_context().params["cohort"]
+            fname = f'{command_name}_{cohort if cohort != "None" else ""}_{DATE}.log'
+        else: 
+            fname = f"{command_name}_{DATE}.log"
+
+        os.makedirs(log_dir, exist_ok=True)
+        
+        level = logging.DEBUG if click.get_current_context().params['verbose'] else logging.INFO
+
+        formatter = daiquiri.formatter.ColorFormatter(fmt=FORMAT)
+        
+        daiquiri.setup(level=level, outputs=(
+            daiquiri.output.Stream(formatter=formatter), 
+            daiquiri.output.File(filename=os.path.join(log_dir, fname), formatter=formatter)
+        ))
+        
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def startup_message(version, initializing_text):
+    
+    author = "Biomedical Genomics Lab - IRB Barcelona"
+    support_email = "stefano.pellegrini@irbbarcelona.org"
+    banner_width = 70
+
+    logger.info("#" * banner_width)
+    logger.info(f"{'#' + ' ' * (banner_width - 2) + '#'}")
+    logger.info(f"{'#' + f'Welcome to Oncodrive3D!'.center(banner_width - 2) + '#'}")
+    logger.info(f"{'#' + ' ' * (banner_width - 2) + '#'}")
+    logger.info(f"{'#' + initializing_text.center(banner_width - 2) + '#'}")
+    logger.info(f"{'#' + ' ' * (banner_width - 2) + '#'}")
+    logger.info(f"{'#' + f'Version: {version}'.center(banner_width - 2) + '#'}")
+    logger.info(f"{'#' + f'Author: {author}'.center(banner_width - 2) + '#'}")
+    logger.info(f"{'#' + f'Support: {support_email}'.center(banner_width - 2) + '#'}")
+    logger.info(f"{'#' + ' ' * (banner_width - 2) + '#'}")
+    logger.info("#" * banner_width)
+    logger.info("")
+
+
+# ===================
+#  Clean and organize
+# ===================
+
+
+def copy_dir(source_dir, destination_dir):
+    """
+    Ccopy the entire directory.
+    """
+
+    logger.debug("Copying directory..")
+    logger.debug(f"From {source_dir}")
+    logger.debug(f"To {destination_dir}")
+    
+    if os.path.exists(source_dir):
+        if os.path.isdir(source_dir):
+            
+            try:
+                shutil.copytree(source_dir, destination_dir)
+                logger.debug("Directory copied successfully!")
+            except Exception as e:
+                logger.warning(f"An error occurred ({e}): Skipping")
+                
+        else:
+            logger.warning("Error while copying directory (source path is not a directory): Skipping")
+    else:
+        logger.warning("Error while copying directory (source path does not exist): Skipping")
+        
+
+def rm_dir(dir_path):             # TO DO: Probably, if the directory is not empty, it should ask for confirmation
+    """
+    Remove directory.
+    """
+    
+    if os.path.isdir(dir_path):
+        shutil.rmtree(dir_path, ignore_errors=True)
+    
+    
+def rm_files(dir_path, ext=[".cif.gz", ".cif"]) -> None:
+    """
+    Remove any file with a given extension in a given directory.
+    """
+    
+    for file in os.listdir(dir_path):
+        file_path = os.path.join(dir_path, file)
+        if any([file.endswith(ext_i) for ext_i in ext]):
+            os.remove(file_path)
+
+
+def clean_directory(path: str, loc: str) -> None:
+    """
+    Clean a directory by removing specific files and subdirectories.
+    """
+
+    if loc == "d":
+        logger.debug(f"Cleaning {path}")
+        rm_files(path, ext=[".tsv", ".csv", ".json", ".txt", ".txt.gz"])
+        dirs = "pae", "pdb_structures", "pdb_structures_mane", "prob_cmaps"
+        path_dirs = [os.path.join(path, d) for d in dirs]
+        for path in path_dirs:
+            rm_dir(path)
+
+    elif loc == "r":
+        # TODO: implement cleaning function for output
+        pass
+
+
+def clean_dir(path: str, loc: str = 'd', txt_file=False) -> None:
+    """
+    Clean it upon request if it already exists.
+    """
+
+    if os.listdir(path) != ['log']:
+        logger.warning(f"Directory {path} already exists and is not empty.")
+
+        overwrite = "y" if click.get_current_context().params['yes'] else input("Clean existing directory? (y/n): ")
+        while overwrite.lower() not in ["y", "yes", "n", "no"]:
+            print("Please choose yes or no")
+            overwrite = input("Clean existing directory? (y/n): ")
+
+        if overwrite.lower() in ["y", "yes"]:
+            clean_directory(path, loc)
+            logger.info(f"Dataset files in {path} have been removed.")
+        else:
+            logger.warning(f"Dataset files in {path} have not been removed.")
+    else:
+        pass
+            
+            
+def clean_temp_files(path: str) -> None:
+    """
+    Clean temp files from dir after completing building the datasets. 
+    """
+    
+    pdb_dir = os.path.join(path, "pdb_structures")
+    rm_files(pdb_dir, ext=[".cif.gz", ".cif", ".tar"])
+    rm_files(os.path.join(path, "pae"), ext=[".json"])
+    rm_dir(os.path.join(path, "pdb_structures_mane"))
+    rm_dir(os.path.join(pdb_dir, "fragmented_pdbs"))