Xue-Jun/StructureBasedSimilarityNetwork

1import hashlib
2import os
3import sys
4from io import StringIO
5from pathlib import Path
6
7import numpy as np
8import pandas as pd
9import rpy2.robjects as ro
10from rpy2.robjects import pandas2ri
11from rpy2.robjects.conversion import localconverter
12
13from r_functions import export_matrix_to_newick_r, export_similarity_network_r
14from usalign_runner import USalignRunner
15
16
17def get_TM_mat_from_df(df):
18    unique_chains = sorted(set(df["#PDBchain1"].unique()).union(set(df["PDBchain2"].unique())))
19    chain_to_idx = {chain: idx for idx, chain in enumerate(unique_chains)}
20    n = len(unique_chains)
21    matrix = np.eye(n)
22    for _, row in df.iterrows():
23        chain1 = row["#PDBchain1"]
24        chain2 = row["PDBchain2"]
25        if chain1 in chain_to_idx and chain2 in chain_to_idx:
26            i = chain_to_idx[chain1]
27            j = chain_to_idx[chain2]
28            matrix[j, i] = row["TM1"]
29            matrix[i, j] = row["TM2"]
30
31    columns_names = [chain.replace("/", "").replace(".pdb:A", "") for chain in unique_chains]
32    df = pd.DataFrame(np.array(matrix), columns=columns_names, index=columns_names)
33    return df
34
35
36def calculate_md5(files):
37    hash_md5 = hashlib.md5()
38    sorted_files = sorted(files, key=lambda x: x.name)
39
40    for file in sorted_files:
41        with open(file.name, "rb") as f:
42            for chunk in iter(lambda: f.read(4096), b""):
43                hash_md5.update(chunk)
44
45    return hash_md5.hexdigest()
46
47
48def save_pdb_files(files, data_dir="./data"):
49    """Save uploaded PDB files to the specified directory."""
50    if not files:
51        return "No files uploaded"
52
53    # Create data directory if it doesn't exist
54    data_path = Path(data_dir)
55    data_path.mkdir(parents=True, exist_ok=True)
56
57    # Calculate MD5 hash for all files
58    md5_hash = calculate_md5(files)
59
60    file_dir = os.path.join(data_path, md5_hash)
61    # file_dir.mkdir(exist_ok=True)
62    try:
63        os.mkdir(file_dir)
64    except Exception:
65        pass
66    file_dir = os.path.join(data_path, md5_hash, "pdb")
67    try:
68        os.mkdir(file_dir)
69    except Exception:
70        pass
71    print(f"Created directory: {file_dir}")
72
73    # Create list file
74    list_file = os.path.join(data_path, md5_hash, "pdb_list")
75
76    filenames = []
77
78    results = []
79    for file in files:
80        # Get original filename
81        original_filename = os.path.basename(file.name)
82        filenames.append(original_filename)
83        # Check if file already exists
84        target_path = os.path.join(file_dir, original_filename)
85        print(f"Saving to: {target_path}")
86
87        # Save the file
88        with open(target_path, "wb") as f:
89            f.write(open(file.name, "rb").read())
90        results.append(f"Saved {original_filename}")
91
92    # Write list file
93    with open(list_file, "w") as f:
94        f.write("\n".join(filenames))
95    results.append(f"Created list file: {list_file}")
96
97    return "\n".join(results)
98
99
100def run_usalign(md5_hash):
101    """Run USalign on the uploaded PDB files and return results as DataFrame."""
102    try:
103        runner = USalignRunner()
104        data_path = Path("./data")
105        pdb_dir = os.path.join(data_path, md5_hash, "pdb")
106        list_file = os.path.join(data_path, md5_hash, "pdb_list")
107        print(str(pdb_dir))
108        print(str(list_file))
109        return_code, stdout, stderr = runner.run_alignment(target_dir=str(pdb_dir), pdb_list_file=str(list_file))
110        print(stdout)
111        print(stderr)
112        if return_code == 0:
113            # Handle potential encoding issues
114            df = pd.read_csv(StringIO(stdout), sep="\t", encoding=sys.getdefaultencoding())
115
116            # Clean up any potential encoding artifacts in column names
117            df.columns = [col.strip() for col in df.columns]
118            return df
119        else:
120            return pd.DataFrame({"Error": [stderr]})
121    except Exception as e:
122        return pd.DataFrame({"Error": [e, stderr]})
123
124
125def run_community_analysis(results_df, data_dir, md5_hash, threshold):
126    """Run community analysis pipeline and return results."""
127    try:
128        # Generate TM matrix
129        tm_matrix = get_TM_mat_from_df(results_df)
130
131        tm_file = os.path.join("data", md5_hash, "tm_matrix.csv")
132        newick_file = os.path.join("data", md5_hash, "clustering.newick")
133        # network_file = os.path.join("data",md5_hash,"network.svg")
134        network_edges_file = os.path.join("data", md5_hash, "network_cytoscape_export.xlsx")
135        # cluster_file = os.path.join("data", md5_hash, "cluster_assignments.csv")
136
137        with localconverter(ro.default_converter + pandas2ri.converter):
138            r_tm_matrix = ro.conversion.py2rpy(tm_matrix)
139
140            result = export_matrix_to_newick_r(r_tm_matrix, newick_file)
141            newick_str = result[0]
142
143            export_similarity_network_r(threshold, r_tm_matrix, network_edges_file)
144
145        # cluster_df.to_csv(cluster_file,index=False)
146        # combined_df.to_csv(network_edges_file,index=False)
147        tm_matrix.to_csv(tm_file)
148        # with open(newick_file, "w") as f:
149        #     f.write(newick_str)
150        # Phylo.write(tree, newick_file, "newick")
151        # fig.savefig(network_file, format="svg", bbox_inches="tight")
152        # plt.close(fig)
153
154        return {
155            "tm_matrix": tm_matrix,
156            "newick_str": newick_str,
157            # "network_fig": fig,
158            "files": [
159                tm_file,
160                newick_file,
161                # network_file,
162                network_edges_file,
163                # cluster_file,
164            ],
165        }
166    except Exception as e:
167        print("Error", str(e))
168        return {"Error": str(e)}
169