source: rotdif/bin/preprocess_dyna_rho.py @ 1547

Last change on this file since 1547 was 1547, checked in by yuexi, 3 years ago

add dyna rho

  • Property svn:executable set to *
File size: 4.9 KB
Line 
1#!/usr/bin/python
2import pandas as pd
3import numpy as np
4import sys
5import re
6def preprocess_out(start_ind,end_ind,all_txt, dat_type):
7    part_data = all_txt[start_ind:end_ind]
8    data_list = []
9    for items in part_data:
10        if len(items) > 0:
11            data_list.append(items.split())
12    if dat_type == "exp":
13        col_name = ["residue","chain","atom1","atom2","freq","rho_exp","rho_pred","rho_error"]
14    elif dat_type == "dyna":
15        col_name = ["residue","chain","atom1","atom2","freq","csa","s2","s2_err","tau_loc","Rex","s2_fast","tau_fast","chi2"]
16    else:
17        print("Unsupported data type %s." % dat_type)
18    output_pd = pd.DataFrame(data = data_list, columns = col_name)
19    resi = list(output_pd["residue"])
20    digi_resi = []
21    #eliminate * in residue number
22    for list_elem in resi:
23        digi_resi.extend(re.findall("\d+", list_elem)) 
24    digi_resi = [int(elem) for elem in digi_resi]
25    output_pd["residue"] = digi_resi
26    sort_output = output_pd.sort_values(['freq', 'residue'], ascending=[True, True]) 
27    freq_set = set(list(sort_output["freq"]))
28    chain_set = set(list(sort_output["chain"]))
29    output_dict = {}
30    output_keys = []
31#no docking involved
32    if len(chain_set) == 1:
33        for items in freq_set:
34            tmp = sort_output[sort_output["freq"]==items]   
35            freq_tag = str(items) + "MHz"
36            output_keys.append(freq_tag)
37            output_dict[freq_tag] = tmp
38        all_rho_exp = []
39        all_rho_pred = []
40        all_rho_err = []
41        all_rho_resi = []
42        if dat_type =="exp":
43            for items in output_keys:
44                rho_resi_list = [float(ele) for ele in list(output_dict[items]["residue"])]   
45                rho_exp_list = [float(ele) for ele in list(output_dict[items]["rho_exp"])]
46                rho_pred_list = [float(ele) for ele in list(output_dict[items]["rho_pred"])]
47                rho_err_list = [float(ele) for ele in list(output_dict[items]["rho_error"])]
48                all_rho_resi.append(rho_resi_list)
49                all_rho_exp.append(rho_exp_list)
50                all_rho_pred.append(rho_pred_list)
51                all_rho_err.append(rho_err_list)
52            all_output = [all_rho_resi, all_rho_exp,all_rho_pred,all_rho_err]
53        elif dat_type =="dyna":
54        #all frequencies have the same dyanmics data
55            resi_list = [float(ele) for ele in list(output_dict[output_keys[0]]["residue"])]
56            s2_list = [float(ele) for ele in list(output_dict[output_keys[0]]["s2"])]
57            tau_loc_list = [float(ele) for ele in list(output_dict[output_keys[0]]["tau_loc"])]
58            rex_list = [float(ele) for ele in list(output_dict[output_keys[0]]["Rex"])]
59            s2_fast_list = [float(ele) for ele in list(output_dict[output_keys[0]]["s2_fast"])]
60            all_output =[resi_list, s2_list, tau_loc_list,rex_list,s2_fast_list]
61    elif len(chain_set) == 2:
62        for items in chain_set:
63            tmp = sort_output[sort_output["chain"]==items]   
64            chain_tag = str(items) + "MHz"
65            output_keys.append(chain_tag)
66            output_dict[chain_tag] = tmp
67        all_s2 = []
68        all_tau_loc = []
69        all_rex = []
70        all_s2_fast = []
71        all_rho_exp = []
72        all_rho_pred = []
73        all_rho_err = []
74        all_rho_resi = []
75        if dat_type =="exp":
76            for items in output_keys:
77                rho_resi_list = [float(ele) for ele in list(output_dict[items]["residue"])]   
78                rho_exp_list = [float(ele) for ele in list(output_dict[items]["rho_exp"])]
79                rho_pred_list = [float(ele) for ele in list(output_dict[items]["rho_pred"])]
80                rho_err_list = [float(ele) for ele in list(output_dict[items]["rho_error"])]
81                all_rho_resi.append(rho_resi_list)
82                all_rho_exp.append(rho_exp_list)
83                all_rho_pred.append(rho_pred_list)
84                all_rho_err.append(rho_err_list)
85            all_output = [all_rho_resi, all_rho_exp,all_rho_pred,all_rho_err]
86        elif dat_type =="dyna":
87        #different chains have different dyanmics data
88            for items in output_keys:
89                resi_list = [float(ele) for ele in list(output_dict[output_keys[0]]["residue"])]
90                s2_list = [float(ele) for ele in list(output_dict[items]["s2"])]
91                tau_loc_list = [float(ele) for ele in list(output_dict[items]["tau_loc"])]
92                rex_list = [float(ele) for ele in list(output_dict[items]["Rex"])]
93                s2_fast_list = [float(ele) for ele in list(output_dict[items]["s2_fast"])]
94                all_rho_resi.append(resi_list)
95                all_s2.append(s2_list)
96                all_tau_loc.append(tau_loc_list)
97                all_rex.append(rex_list)
98                all_s2_fast.append(s2_fast_list)
99            all_output =[all_rho_resi, all_s2,all_tau_loc,all_rex,all_s2_fast]
100       
101    else:
102        sys.stderr.write("More than 2 chains or fewer than 1 chain are not supported\n")
103    return all_output
Note: See TracBrowser for help on using the repository browser.