Oncodrive3D 1.0.4__py3-none-any.whl

scripts/plotting/build_annotations.py

@@ -0,0 +1,102 @@
+import logging
+import os
+import pandas as pd
+
+import daiquiri
+
+from scripts import __logger_name__
+from scripts.globals import copy_dir
+from scripts.plotting.utils import clean_annot_dir
+from scripts.plotting.utils import get_species
+from scripts.plotting.stability_change import download_stability_change, parse_ddg_rasp
+from scripts.plotting.pdb_tool import run_pdb_tool, parse_pdb_tool
+from scripts.plotting.pfam import get_pfam
+from scripts.plotting.uniprot_feat import get_uniprot_feat
+
+logger = daiquiri.getLogger(__logger_name__ + ".plotting.build")
+
+logging.getLogger('urllib3.connectionpool').setLevel(logging.WARNING)
+
+
+
+# TODO: fix multiprocessing on DDG
+# TODO: multiprocessing on pdb_tool
+
+# =================
+# BUILD ANNOTATIONS 
+# ================= 
+
+def get_annotations(data_dir,
+                    output_dir,
+                    ddg_dir,
+                    #path_pdb_tool_sif,
+                    organism,
+                    cores):
+    """
+    Main function to build annotations to generate annotated plots.
+    """
+
+    # Empty directory and load sequence df
+    clean_annot_dir(output_dir, 'd')
+
+    # # Get ddG
+    species = get_species(organism)
+    logger.info(f"Obtaining stability change..")
+    if species == "Homo sapiens" or species == "Mus musculus":
+        ddg_output = os.path.join(output_dir, "stability_change")
+        os.makedirs(ddg_output, exist_ok=True)
+        if ddg_dir is not None:
+            # Copy ddG from path
+            temp_ddg_path = os.path.join(output_dir, "stability_change_temp")
+            copy_dir(source_dir=ddg_dir, destination_dir=temp_ddg_path)
+        else:    
+            # Download ddG
+            temp_ddg_path = download_stability_change(ddg_output, cores)
+        logger.info(f"Completed!")
+        
+        ## TODO: Optimize DDG parsing 
+        ##       - only one protein is allocated to one process every time
+        ##       - a list of proteins should be allocated instead
+        
+        # Parsing DDG
+        logger.info(f"Parsing stability change..")
+        parse_ddg_rasp(temp_ddg_path, ddg_output, cores)
+        logger.info(f"Parsing completed!")
+    else:
+        logger.warning(f"Currently, stability change annotation is not available for {species} but only for Homo sapiens: Skipping...")
+    
+    ## TODO: Enable multiprocessing for PDB_Tool
+    ## TODO: Evaluate the possibility of installing PDB_Tool instead of using container
+    
+    # Run PDB_Tool
+    logger.info(f"Extracting pACC and 2° structure..")
+    path_pdb_structure = os.path.join(data_dir, "pdb_structures")
+    pdb_tool_output = run_pdb_tool(input_dir=path_pdb_structure, output_dir=output_dir)
+    logger.info(f"Extraction completed!")
+    
+    # Parse PDB_Tool
+    logger.info(f"Parsing pACC and 2° structures..")
+    parse_pdb_tool(input_dir=pdb_tool_output, output_dir=output_dir)
+    logger.info(f"Parsing completed!")
+    
+    # Get Pfam annotations
+    logger.info(f"Downloading and parsing Pfam..")
+    seq_df = pd.read_table(os.path.join(data_dir, "seq_for_mut_prob.tsv"))    
+    pfam_df = get_pfam(seq_df = seq_df, 
+                       output_tsv = os.path.join(output_dir, "pfam.tsv"),
+                       organism = species)
+    logger.info(f"Completed!")
+    
+    # Get Uniprot features
+    logger.info(f"Downloading and parsing Features from Uniprot..")
+    get_uniprot_feat(seq_df = seq_df, 
+                     pfam_df = pfam_df,
+                     output_tsv = os.path.join(output_dir, "uniprot_feat.tsv"))      
+    logger.info(f"Completed!")
+    
+
+if __name__ == "__main__":
+    get_annotations(data_dir = "/workspace/nobackup/scratch/oncodrive3d/datasets",
+                    output_dir = "/workspace/nobackup/scratch/oncodrive3d/annotations",
+                    organism = "Homo sapiens",
+                    cores = 4)

scripts/plotting/chimerax_plot.py

@@ -0,0 +1,251 @@
+"""
+Use ChimeraX to generate 3D structures colored by metrics
+"""
+
+import pandas as pd   
+import numpy as np
+import os 
+import subprocess
+import math
+import daiquiri
+
+from scripts import __logger_name__
+logger = daiquiri.getLogger(__logger_name__ + ".plotting.chimerax_plot")
+
+
+def create_attribute_file(path_to_file, 
+                        df, 
+                        attribute_col,
+                        pos_col="Pos",
+                        attribute_name="local_attribute",
+                        gene="Gene_name", 
+                        protein="Protein_name"):
+
+    logger.info(f"Saving {path_to_file}")
+    with open(path_to_file, 'w') as f:
+        f.write("#\n")
+        f.write(f"#  Mutations in volume for {protein} ({gene})\n")
+        f.write("#\n")
+        f.write("#  Use this file to assign the attribute in Chimera with the \n")
+        f.write("#  Define Attribute tool or the command defattr.\n")
+        f.write("#\n")
+        f.write(f"attribute: {attribute_name}\n")
+        f.write("recipient: residues\n")
+
+        for _, row in df.iterrows():
+            f.write(f"\t:{int(row[pos_col])}\t{row[attribute_col]}\n")
+            
+            
+def round_to_first_nonzero(num):
+    if num == 0:
+        return 0 
+    
+    scale = -int(math.floor(math.log10(abs(num))))
+    shifted_num = num * (10 ** scale)
+    rounded_num = round(shifted_num)
+    result = rounded_num / (10 ** scale)
+    
+    return result
+            
+
+def is_float_an_integer(value):
+    if isinstance(value, float):
+        return value.is_integer()
+    return False
+
+
+def get_intervals(attribute_vector, attribute):
+    
+    max_value = attribute_vector.max()
+    min_value = attribute_vector.min()
+    min_value = 0 if attribute == "score" else 1
+    intervals = np.linspace(min_value, max_value, 5)
+    intervals = np.round(intervals, 2) if attribute == "score" else np.round(intervals)
+    if attribute == "logscore":
+        pos_values = np.linspace(0, max_value, 3)
+        intervals = np.round([-max_value, -pos_values[1], 0, pos_values[1], max_value], 2)
+    intervals = [round(n) if is_float_an_integer(n) else n for n in intervals]
+        
+    return intervals
+
+
+def get_palette(intervals, type="diverging"):
+    
+    # Diverging palette
+    if type == "diverging":
+        return f"{intervals[0]},#0571B0:{intervals[1]},#92C5DE:{intervals[2]},white:{intervals[3]},#F4A582:{intervals[4]},#CA0020"
+    
+    # Sequential palette
+    else:
+        return f"{intervals[0]},#FFFFB2:{intervals[1]},#FECC5C:{intervals[2]},#FD8D3C:{intervals[3]},#F03B20:{intervals[4]},#BD0026"
+
+
+def get_chimerax_command(chimerax_bin, 
+                         pdb_path, 
+                         chimera_output_path, 
+                         attr_file_path, 
+                         attribute, 
+                         intervals, 
+                         gene, 
+                         uni_id,
+                         labels,
+                         i,
+                         f,
+                         cohort="",
+                         clusters=None,
+                         pixelsize=0.1,
+                         transparent_bg=False):
+    
+    palette = get_palette(intervals, type="diverging") if attribute == "logscore" else get_palette(intervals, type="sequential")
+    transparent_bg = " transparentBackground  true" if transparent_bg else ""
+    
+    chimerax_command = (
+        f"{chimerax_bin} --nogui --offscreen --silent --cmd "
+        f"\"open {pdb_path}; "
+        "set bgColor white; "
+        "color lightgray; "
+        f"open {attr_file_path}; "
+        f"color byattribute {attribute} palette {palette}; "
+        f"key {palette} :{intervals[0]} :{intervals[1]} :{intervals[2]} :{intervals[3]} :{intervals[4]} pos 0.35,0.03 fontSize 4 size 0.3,0.02;"
+        f"2dlabels create label text '{labels[attribute]}' size 6 color darkred xpos 0.34 ypos 0.065;"
+        f"2dlabels create title text '{gene} - {uni_id}-F{f} ' size 6 color darkred xpos 0.35 ypos 0.93;"
+        "hide atoms;"
+        "show cartoons;"
+        "lighting soft;"
+        "graphics silhouettes true width 1.3;"
+        "zoom;"
+    )
+    
+    if clusters is not None and len(clusters) > 0:
+        for pos in clusters:
+            chimerax_command += f"marker #10 position :{pos} color #dacae961 radius 5.919;"
+        cluster_tag = "_clusters"
+    else:
+        cluster_tag = ""
+    
+    output_path = os.path.join(chimera_output_path, f"{cohort}_{i}_{gene}_{attribute}{cluster_tag}.png")
+    chimerax_command += f"save {output_path} pixelSize {pixelsize} supersample 3{transparent_bg};exit\""
+    
+    return chimerax_command
+            
+
+def generate_chimerax_plot(output_dir,
+                            gene_result_path,
+                            pos_result_path,
+                            datasets_dir,
+                            seq_df_path,
+                            cohort,
+                            max_genes,
+                            pixel_size,
+                            cluster_ext,
+                            fragmented_proteins,
+                            transparent_bg,
+                            chimerax_bin):
+
+    seq_df = pd.read_csv(seq_df_path, sep="\t")
+    gene_result = pd.read_csv(gene_result_path)
+    result = pd.read_csv(pos_result_path)
+    if "Ratio_obs_sim" in result.columns:
+        result = result.rename(columns={"Ratio_obs_sim" : "Score_obs_sim"})
+    result["Logscore_obs_sim"] = np.log(result["Score_obs_sim"])
+
+    # Process each gene
+    genes = gene_result[gene_result["C_gene"] == 1].Gene.unique()
+    if len(genes) > 0:
+        
+        chimera_out_path = os.path.join(output_dir, f"{cohort}.chimerax")
+        chimera_attr_path = os.path.join(chimera_out_path, "attributes")
+        chimera_plots_path = os.path.join(chimera_out_path, "plots")
+        for path in [chimera_out_path, chimera_attr_path, chimera_plots_path]:
+            os.makedirs(path, exist_ok=True)
+                
+        for i, gene in enumerate(genes[:max_genes]):
+            logger.info(f"Processing {gene}")
+            
+            # Attribute files
+            logger.debug("Preprocessing for attribute files..")
+            result_gene = result[result["Gene"] == gene]
+            if cluster_ext:
+                clusters = result_gene[result_gene.C == 1].Pos.values
+            else:
+                clusters = result_gene[(result_gene.C == 1) & (result_gene.C_ext == 0)].Pos.values            
+            len_gene = len(seq_df[seq_df["Gene"] == gene].Seq.values[0])
+            result_gene = pd.DataFrame({"Pos" : range(1, len_gene+1)}).merge(
+                result_gene[["Pos", 
+                             "Mut_in_res", 
+                             "Mut_in_vol", 
+                             "Score_obs_sim", 
+                             "Logscore_obs_sim"]], on="Pos", how="left")
+
+            uni_id, f = seq_df[seq_df["Gene"] == gene][["Uniprot_ID", "F"]].values[0]
+            pdb_path = os.path.join(datasets_dir, "pdb_structures", f"AF-{uni_id}-F{f}-model_v4.pdb")
+            
+            labels = {"mutres" : "Mutations in residue ", 
+                      "mutvol" : "Mutations in volume ",
+                      "score" : "   Clustering score ",
+                      "logscore" : "log(Clustering score) "}
+            cols = {"mutres" : "Mut_in_res", 
+                    "mutvol" : "Mut_in_vol",
+                    "score" : "Score_obs_sim",
+                    "logscore" : "Logscore_obs_sim"}                                            
+            
+            if fragmented_proteins == False:
+                if f != 1:
+                    logger.debug(f"Fragmented protein processing {fragmented_proteins}: Skipping {gene} ({uni_id}-F{f})..")
+                    continue
+                
+            if os.path.exists(pdb_path):
+                
+                for attribute in ["mutres", "mutvol", "score", "logscore"]:
+                    
+                    logger.debug("Generating attribute files..")
+                    attr_file_path = f"{chimera_attr_path}/{gene}_{attribute}.defattr"
+                    create_attribute_file(path_to_file=attr_file_path,
+                                        df=result_gene.dropna(),
+                                        attribute_col=cols[attribute],
+                                        attribute_name=attribute)
+                    
+                    attribute_vector = result_gene[cols[attribute]]
+                    intervals = get_intervals(attribute_vector, attribute)
+ 
+                    logger.debug("Generating 3D images..")
+                    chimerax_command = get_chimerax_command(chimerax_bin, 
+                                                            pdb_path, 
+                                                            chimera_plots_path, 
+                                                            attr_file_path, 
+                                                            attribute, 
+                                                            intervals, 
+                                                            gene,
+                                                            uni_id,
+                                                            labels,
+                                                            i,
+                                                            f,
+                                                            cohort,
+                                                            pixelsize=pixel_size,
+                                                            transparent_bg=transparent_bg)
+                    subprocess.run(chimerax_command, shell=True)
+                    logger.debug(chimerax_command)
+                    
+                    if attribute == "score" or attribute == "logscore":
+                        chimerax_command = get_chimerax_command(chimerax_bin, 
+                                                                pdb_path, 
+                                                                chimera_plots_path, 
+                                                                attr_file_path, 
+                                                                attribute, 
+                                                                intervals, 
+                                                                gene,
+                                                                uni_id,
+                                                                labels,
+                                                                i,
+                                                                f,
+                                                                cohort,
+                                                                clusters=clusters,
+                                                                pixelsize=pixel_size,
+                                                                transparent_bg=transparent_bg)
+                        subprocess.run(chimerax_command, shell=True)
+                        logger.debug(chimerax_command)
+                        
+            else:
+                logger.warning(f"PDB path missing: {pdb_path}")
+    else:
+        logger.info("Nothing to plot!")

scripts/plotting/pdb_tool.py

@@ -0,0 +1,149 @@
+import logging
+import os
+
+import daiquiri
+import subprocess
+import pandas as pd
+import numpy as np
+import re
+import shutil
+from tqdm import tqdm
+
+from scripts import __logger_name__
+
+logger = daiquiri.getLogger(__logger_name__ + ".plotting.pdb_tool")
+
+logging.getLogger('urllib3.connectionpool').setLevel(logging.WARNING)
+
+
+
+# ========
+# PDB_Tool
+# ========
+
+
+def decompress_pdb_gz(input_dir):
+    """
+    Extract all .gz PDB files in directory.
+    """
+    
+    files_in_directory = os.listdir(input_dir)
+    gz_file_present = any(file.endswith(".pdb.gz") for file in files_in_directory)
+    if gz_file_present:
+        logger.debug("Decompressing .gz PDB files...")
+        command = f'gunzip -q {os.path.join(input_dir, "*.pdb.gz")}'
+        subprocess.run(command, shell=True)
+        logger.debug("Decompression complete!")
+    
+
+def run_pdb_tool(input_dir, output_dir, f="4"):
+    """
+    Use PDB_Tool to extract features from all pdb files in directory.
+    """
+
+    pdb_tool_output = os.path.join(output_dir, "pdb_tool")
+    if not os.path.isdir(pdb_tool_output):
+        os.makedirs(pdb_tool_output)
+        logger.debug(f'mkdir {pdb_tool_output}')
+    
+    decompress_pdb_gz(input_dir)
+    pdb_files = [file for file in os.listdir(input_dir) if file.endswith(".pdb")]    
+    logger.debug("Running PDB_Tool...")
+    for file in tqdm(pdb_files, desc="Running PDB_Tool"):
+        output = os.path.join(pdb_tool_output, file.replace(".pdb", ".feature"))
+        # Singularity container
+        #subprocess.run(["singularity", "exec", f"{pdb_tool_sif_path}", "/PDB_Tool/PDB_Tool", "-i", f"{input_dir}/{file}", "-o", output, "-F", f])            
+        # Added to $PATH
+        subprocess.run(["PDB_Tool", "-i", os.path.join(input_dir, file), "-o", output, "-F", f])   
+        
+    return pdb_tool_output
+            
+            
+def load_pdb_tool_file(path):
+    """
+    Parse .feature file obtained by PDB_Tool.
+    """
+    
+    with open(path, "r") as f:
+        lines = f.readlines()
+        lst_lines = []
+        for l in lines:
+            lst_lines.append(l.strip().split())
+   # return lst_lines
+    df = pd.DataFrame(lst_lines[1:], columns = lst_lines[0]).iloc[:,:9]
+    df = df.drop("Missing", axis=1)
+    df = df.rename(columns={"Num" : "Pos"})
+
+    for c in df.columns:
+        if c == "Pos" or c == "ACC" or c == "CNa" or c == "CNb":
+            data_type = int
+        else:
+            data_type = float
+        try:
+            df[c] = df[c].astype(data_type)
+        except:
+            pass
+    return df
+
+
+def get_pdb_tool_file_in_dir(path):
+    """
+    Get the list of PDB_Tool .feature files from directory.
+    """
+    
+    list_files = os.listdir(path)
+    ix = [re.search('.\.feature$', x) is not None for x in list_files]
+    
+    return list(np.array(list_files)[np.array(ix)])
+
+
+def load_all_pdb_tool_files(path):
+    """
+    Get the features of all proteins in the directory.
+    """
+    
+    df_list = []
+    feature_file_list = get_pdb_tool_file_in_dir(path)
+    
+    for file in tqdm(feature_file_list, desc="Parsing PDB_tool output"):
+        df = load_pdb_tool_file(os.path.join(path, file))
+        identifier = file.split("-")
+        df["Uniprot_ID"] = identifier[1]
+        df["F"] = identifier[2].replace("F", "")
+        df_list.append(df)
+
+    return pd.concat(df_list).reset_index(drop=True)
+
+
+def pdb_tool_to_3s_sse(df):
+    """
+    Reduce secondary structure from 8 to 3 states.
+    """
+
+    mapper = {"H":"Helix", 
+              "G":"Helix", 
+              "I":"Helix", 
+              "L":"Coil", 
+              "T":"Coil", 
+              "S":"Coil", 
+              "B":"Coil", 
+              "E":"Ladder"}
+    df["SSE"] = df["SSE"].map(mapper)
+
+    return df
+
+
+def parse_pdb_tool(input_dir : str, output_dir : str):
+    """
+    Parse PDB_Tool .feature files inclued in the path into a unique df.
+    """
+    
+    pdb_tool_df = load_all_pdb_tool_files(input_dir)
+    pdb_tool_df = pdb_tool_to_3s_sse(pdb_tool_df)
+    pdb_tool_df = pdb_tool_df.drop(columns=["CLE", "ACC", "CNa", "CNb"])
+    pdb_tool_df.to_csv(os.path.join(output_dir, "pdb_tool_df.tsv"), sep="\t", index=False)
+    try:
+        logger.debug(f"Deleting {input_dir}")
+        shutil.rmtree(input_dir)
+    except OSError as e:
+        logger.debug(f"Cold not delete {input_dir}\nError: {e}")

scripts/plotting/pfam.py

@@ -0,0 +1,94 @@
+import logging
+import os
+import subprocess
+
+import daiquiri
+import pandas as pd
+
+from scripts import __logger_name__
+
+logger = daiquiri.getLogger(__logger_name__ + ".plotting.pfam")
+
+logging.getLogger('urllib3.connectionpool').setLevel(logging.WARNING)
+
+
+
+def add_pfam_metadata(pfam, seq_df):
+    """
+    Add Ensembl transcript and gene info and rename cols to
+    be merged with Uniprot features dataframe.
+    """
+
+    # Add metadata to PFAM
+    pfam = seq_df[["Gene", "Uniprot_ID", "Ens_Transcr_ID", "Ens_Gene_ID"]].merge(
+        pfam, how="left", on=["Ens_Transcr_ID", "Ens_Gene_ID"])
+    pfam = pfam.dropna(how="all", subset=["Pfam_start", "Pfam_end"]).reset_index(drop=True)
+
+    # Prepare to merge
+    pfam["Type"] = "DOMAIN"
+    pfam["Evidence"] = "Pfam"
+    pfam = pfam.rename(columns={"Pfam_start" : "Begin",
+                                "Pfam_end" : "End",
+                                "Pfam_name" : "Description",
+                                "Pfam_description" : "Full_description"})
+
+    return pfam
+
+
+def get_pfam(seq_df, output_tsv, organism):
+    """
+    Download and parse Pfam coordinates, name, description,
+    and Pfam ID to Transcript ID mapping.
+    """
+
+    status = "INIT"
+    i = 0
+    if organism == "Homo sapiens":
+        ensembl_gene_dataset = "hsapiens_gene_ensembl"
+    elif organism == "Mus musculus":
+        ensembl_gene_dataset = "mmusculus_gene_ensembl"
+    else:
+        logger.error(f"Invalid organism: {organism}. Expected 'Homo sapiens' or 'Mus musculus'.")
+        raise ValueError(f"Invalid organism: {organism}. Must be 'Homo sapiens' or 'Mus musculus'.")
+
+    while status != "PASS":
+        if i < 5:
+            try:
+                # Pfam coordinates
+                logger.debug("Downloading and parsing Pfam coordinates...")
+                url_query = 'http://jan2024.archive.ensembl.org/biomart/martservice?query='
+                query = f'<?xml version="1.0" encoding="UTF-8"?><!DOCTYPE Query><Query  virtualSchemaName = "default" formatter = "TSV" header = "0" uniqueRows = "0" count = "" datasetConfigVersion = "0.6" ><Dataset name = "{ensembl_gene_dataset}" interface = "default" ><Attribute name = "ensembl_gene_id" /><Attribute name = "ensembl_transcript_id" /><Attribute name = "pfam_start" /><Attribute name = "pfam_end" /><Attribute name = "pfam" /></Dataset></Query>'
+                url = url_query + query
+                command = [f"wget", "-q", "-O", "pfam_coordinates.tsv", url]
+                subprocess.run(command)
+                pfam = pd.read_csv("pfam_coordinates.tsv", sep="\t", header=None)
+                pfam.columns = ["Ens_Gene_ID", "Ens_Transcr_ID", "Pfam_start", "Pfam_end", "Pfam_ID"]
+
+                # ID database
+                logger.debug("Downloading and parsing Pfam ID database...")
+                url = "https://ftp.ebi.ac.uk/pub/databases/Pfam/current_release/database_files/pfamA.txt.gz"
+                command = ["wget", "-q", "-O", "pfam_id.tsv.gz", url]
+                subprocess.run(command)
+                pfam_id = pd.read_csv("pfam_id.tsv.gz", compression='gzip', sep='\t', header=None).iloc[:,[0,1,3]]
+                pfam_id.columns = "Pfam_ID", "Pfam_name", "Pfam_description"
+
+                # Merge and save
+                pfam = pfam.merge(pfam_id, how="left", on="Pfam_ID")
+                pfam = pfam.dropna(how="all", subset=["Pfam_start", "Pfam_end"]).reset_index(drop=True)
+                pfam = add_pfam_metadata(pfam, seq_df)
+                pfam.to_csv(output_tsv, index=False, sep="\t")
+
+                # Delete temp files
+                os.remove("pfam_coordinates.tsv")
+                os.remove("pfam_id.tsv.gz")
+                status = "PASS"
+
+                return pfam
+
+            except Exception as e:
+                status = "FAIL"
+                logger.warning(f"Error while downloading Pfam: {e}")
+                logger.warning("Retrying download...")
+                i += 1
+        else:
+            logger.error(f'Download Pfam: {status}')