V-ToCs/Runalign_tm_rms_multiproc.py at main · DigBioLab/V-ToCs · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import subprocess
import os
import re
import pandas as pd
import numpy as np
from concurrent.futures import ProcessPoolExecutor, as_completed


def execute_command(command_full, count, total_count):
    command, i_protein, target = command_full
    res = subprocess.run(command, stdout=subprocess.PIPE)
    if res.returncode == 0:
        try:
            out = res.stdout.decode()
            TMscore = re.findall(r'TM-score= (\d*\.?\d*)',out)[1]
            TMscore = float(TMscore)
            TMdist = 1 - TMscore
            rmsd = re.findall(r'RMSD=\s+(\d*\.?\d*)',out)[0]
        except IndexError:
            TMscore = 'Failed_re'
            TMdist =  'Failed_re'
            rmsd =    'Failed_re'
            print(f'Failed {target} and {i_protein}, TMalign')

        # follow progress
        if count % 300 == 0:
            print(f'{count/total_count*100:.2f}%; TMscore = {TMscore}')
            with open('runalign_log_processpool.txt', 'a') as f:
                f.write(f'{count/total_count*100:.2f}%; TMscore = {TMscore}\n')
        return (TMscore, TMdist, rmsd, i_protein, target, out)
    else:
        raise ValueError(f'execute_command failed: i = {i_protein}, target = {target}')
        #return None


def Runalign_tm_rms():
    path = 'data'
    # files to analyze
    folder = f'{path}/Pdb_files/modified'
    files = []
    for file in os.listdir(folder):
        if file.endswith('pdb'):
            files.append(file)

    files = sorted(files)

    df = pd.DataFrame(np.nan, columns=files, index=files)           # original TM scores
    df_tm_dist = pd.DataFrame(np.nan, columns=files, index=files)   # distance scores from TM
    df_rms = pd.DataFrame(np.nan, columns=files, index=files)       # RMSD scores

    for i_df in [df, df_tm_dist, df_rms]:
        print(len(i_df), len(i_df.columns))


    c = 0   # count
    prog = 0
    targets = files.copy()

    print('Initialized')

    command_list = []

    for t in files:
        target = targets[0]
        df.loc[target][target] = 1
        df_tm_dist.loc[target][target] = 0
        df_rms[target][target] = 0

        if len(targets) > 1:

            for i in targets[1:]:
                #print(f'Comparing {i}, {target}')
                command = f'./TMalign {folder}/{i} {folder}/{target}'.split(' ')
                # argument with the command and with the proteins involved
                command_full = [command, i, target]
                command_list.append(command_full)
            tpop = targets.pop(0)
            #print(targets,tpop)

        else:
            print('Length one, finished')
        prog += 1
        print(f'{prog} out of {len(files)}')

    total_count = len(command_list)
    print(f'{total_count} comparisons')

    with ProcessPoolExecutor() as executor:
        futures = [executor.submit(execute_command,
                                   command_full, count, total_count) for
                                   (count, command_full) in enumerate(command_list)]


        for future in as_completed(futures):
            output = future.result()

            # save scores
            if output is not None:
                TMscore, TMdist, rmsd, i, target, out = output
                # save TMalign outputs
                '''with open('pylog.txt','a') as fh:
                    #fh.write(out)
                    #fh.write('\n')'''

                # Original TM scores
                df.loc[i][target] = float(TMscore)
                df.loc[target][i] = float(TMscore)

                # Distance scores from TM
                df_tm_dist.loc[i][target] = float(TMdist)
                df_tm_dist.loc[target][i] = float(TMdist)

                # RMSD scores
                df_rms.loc[i][target] = float(rmsd)
                df_rms.loc[target][i] = float(rmsd)

            else:
                raise ValueError(f'execute_command failed: i = {i}, target = {target}')


    print('Cleaning names')

    df.columns = [i[:-17] if 'sel' in i else i[:-4] for i in df.columns]
    df.index = [i[:-17] if 'sel' in i else i[:-4] for i in df.index]

    df_tm_dist.columns = [i[:-17] if 'sel' in i else i[:-4] for i in df_tm_dist.columns]
    df_tm_dist.index = [i[:-17] if 'sel' in i else i[:-4] for i in df_tm_dist.index]

    df_rms.columns = [i[:-17] if 'sel' in i else i[:-4] for i in df_rms.columns]
    df_rms.index = [i[:-17] if 'sel' in i else i[:-4] for i in df_rms.index]

    print('Writing dataframes')
    df.to_csv(f'{path}/Dataframe_TM.csv', sep='\t')
    df_tm_dist.to_csv(f'{path}/Dataframe_TMdist.csv', sep='\t')
    df_rms.to_csv(f'{path}/Dataframe_RMS.csv', sep='\t')

    print('Finished', c)


if __name__ == "__main__":
    Runalign_tm_rms()