source: rotdif/bin/plot_exp_elmdock.py @ 1548

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

add plot exp elmdock

  • Property svn:executable set to *
File size: 10.0 KB
Line 
1#!usr/bin/python
2import re
3import pandas as pd
4import numpy
5import os
6#import our functions
7from elmdock_split import split_data
8
9def exp_plot(color_list, dyna_flag, elm_flag, elmdock_flag):
10    iso_output = split_data(dyna_flag, elm_flag, elmdock_flag)[0]
11    for root, dirs, files in os.walk('./'):
12        for name in files:
13            filename = os.path.join(root, name)
14            if '.txt' in filename:
15                exp_file = filename
16    ##preprocess exp_file in multi-frequencies
17    with open(exp_file) as in_f:
18        tmp_lists = in_f.read().splitlines()
19        exp_list = []
20        for items in tmp_lists:
21            if len(items) > 0:
22                exp_list.append(items.split()) 
23    exp_pd = pd.DataFrame(data = exp_list[1:],\
24columns = ["residue","chain","atom 1","atom 2","magnet","T1","T1 error","T2","T2 error","NOE","NOE error"])
25    resi = list(exp_pd["residue"])
26    digi_resi = []
27    #eliminate * in residue number
28    for list_elem in resi:
29        digi_resi.extend(re.findall("\d+", list_elem)) 
30    digi_resi = [int(elem) for elem in digi_resi]
31    exp_pd["residue"] = digi_resi
32    sort_exp = exp_pd.sort_values(['magnet', 'residue'], ascending=[True, True])
33    freq_set = set(list(exp_pd["magnet"]))
34    chain_set = set(list(exp_pd["chain"]))
35    exp_dict = {} 
36    #for multi-field, it should be frequencies; for multi-chains, it should be chains.
37    exp_keys = []
38    #split data based on frequency difference
39    if len(chain_set) == 0:
40        sys.stderr.write("No chain information in the data!")
41    elif len(chain_set) == 1:   
42        for items in freq_set:
43            #tmp = sort_exp[sort_exp["magnet"]==items]
44            tmp = exp_pd[exp_pd["magnet"]==items]
45            exp_dict[items] = tmp
46            freq_tag = str(items) + "MHz"
47            exp_keys.append(freq_tag)
48#from previous block, iso, axi, ani all have the same exp data
49        rho_exp = iso_output[0]
50        rho_err = iso_output[2]
51        rho_block = []
52#rho data has the same order as rho exp
53        for i in range(len(exp_keys)):
54            rho_data ={
55         "x" : list(range(len(rho_exp[i]))),
56         "y" : rho_exp[i],
57         "error_y": {
58            "array" : rho_err[i],
59            "visible" : True 
60         }, 
61         "xaxis": "x4",
62         "yaxis": "y4", 
63         "mode" : "lines+markers",
64         "line": {"color": color_list[i]},
65         "marker": {"color": color_list[i]},
66         "name" : exp_keys[i],
67         "legendgroup": str(i),
68         "showlegend" : True
69       }   
70            rho_block.append(rho_data)
71        T1_block = []
72        T2_block = []
73        NOE_block = []
74        j = 0
75        for items in freq_set:
76            T1_list = [float(ele) for ele in list(exp_dict[items]["T1"])]
77            T1_err_list = [float(ele) for ele in list(exp_dict[items]["T1 error"])]
78            T2_list = [float(ele) for ele in list(exp_dict[items]["T2"])]
79            T2_err_list = [float(ele) for ele in list(exp_dict[items]["T2 error"])]
80            NOE_list = [float(ele) for ele in list(exp_dict[items]["NOE"])]
81            NOE_err_list = [float(ele) for ele in list(exp_dict[items]["NOE error"])]
82        #for legend
83            T1_data ={
84         "x" : list(exp_dict[items]["residue"]),
85         "y" : T1_list,
86         "error_y": {
87            "array" : T1_err_list,
88            "visible" : True
89         },
90         "xaxis": "x",
91         "yaxis": "y",
92         "mode" : "lines+markers",
93         "line": {"color": color_list[j]},
94         "marker": {"color": color_list[j]},
95         "name" : items,
96         "legendgroup" : str(j),
97         "showlegend" : False
98       }
99            T1_block.append(T1_data)
100
101            T2_data ={
102         "x" : list(exp_dict[items]["residue"]),
103         "y" : T2_list,
104         "error_y": {
105            "array" : T2_err_list,
106            "visible" : True 
107         },
108         "xaxis": "x3",
109         "yaxis": "y3",   
110         "mode" : "lines+markers",
111         "line": {"color": color_list[j]},
112         "marker": {"color": color_list[j]},
113         "name" : items,
114         "legendgroup": str(j),
115         "showlegend" : False
116       }
117            T2_block.append(T2_data)
118
119            NOE_data ={
120         "x" : list(exp_dict[items]["residue"]),
121         "y" : NOE_list,
122         "error_y": {
123            "array" : NOE_err_list,
124            "visible" : True
125         },
126         "xaxis": "x2",
127         "yaxis": "y2",
128         "mode" : "lines+markers",
129         "line": {"color": color_list[j]},
130         "marker": {"color": color_list[j]},
131         "name" : items,
132         "legendgroup" : str(j),
133         "showlegend" : False
134       }
135            NOE_block.append(NOE_data)
136            j += 1
137
138#flat list of lists
139        list_of_list = [rho_block,T1_block,T2_block,NOE_block]
140        list_of_exp = [item for sublist in list_of_list for item in sublist]
141        exp_plotly = {
142   "data" : list_of_exp,
143   "layout" : {
144       "title" : "Experimental Data",
145       "xaxis" : {
146         "title" : "residue number",
147         "domain" : [0,0.45]
148       },
149       "yaxis" : {
150         "title" : "R1_exp",
151         "domain" : [0.55,1]
152       },
153       "xaxis2" : {
154         "title" : "residue number",
155         "domain": [0.55,1],
156         "anchor": "y2"
157       },
158       "yaxis2" : {
159         "title" : "NOE_exp",
160         "domain": [0,0.35],
161         "anchor": "x2"
162       },
163       "xaxis3" : {
164         "title": "residue number",
165         "domain": [0,0.45],
166         "anchor": "y3"
167       },
168       "yaxis3" : {
169         "title": "R2_exp",
170         "domain": [0,0.35],
171         "anchor": "x3"
172       },
173       "xaxis4" : {
174         "title": "residue number",
175         "domain": [0.55,1],
176         "anchor": "y4"
177       },
178       "yaxis4" : {
179         "title": "Rho_exp",
180         "domain": [0.65,1],
181         "anchor": "x4"
182       }
183   }
184}
185    elif len(chain_set) == 2:
186        for items in chain_set:
187            chain_tag = "chain " + str(items)
188            exp_keys.append(chain_tag)
189            tmp = sort_exp[sort_exp["chain"]==items]
190            exp_dict[items] = tmp
191#from previous block, iso, axi, ani all have the same exp data
192        rho_exp = iso_output[0]
193        rho_err = iso_output[2]
194        rho_block = []
195#rho data has the same order as rho exp
196        for i in range(len(exp_keys)):
197            rho_data ={
198         "x" : list(range(len(rho_exp[i]))),
199         "y" : rho_exp[i],
200         "error_y": {
201            "array" : rho_err[i],
202            "visible" : True 
203         },   
204         "xaxis":"x4",
205         "yaxis":"y4",
206         "mode" : "lines+markers",
207         "line": {"color": color_list[i]},
208         "marker": {"color": color_list[i]},
209         "name" : exp_keys[i],
210         "legendgroup": str(i),
211         "showlegend" : True
212       }   
213            rho_block.append(rho_data)
214        T1_block = []
215        T2_block = []
216        NOE_block = []
217        j = 0 
218        for items in chain_set:
219            T1_list = [float(ele) for ele in list(exp_dict[items]["T1"])]
220            T1_err_list = [float(ele) for ele in list(exp_dict[items]["T1 error"])]
221            T2_list = [float(ele) for ele in list(exp_dict[items]["T2"])]
222            T2_err_list = [float(ele) for ele in list(exp_dict[items]["T2 error"])]
223            NOE_list = [float(ele) for ele in list(exp_dict[items]["NOE"])]
224            NOE_err_list = [float(ele) for ele in list(exp_dict[items]["NOE error"])]
225        #for legend
226            T1_data ={
227         "x" : list(exp_dict[items]["residue"]),
228         "y" : T1_list,
229         "error_y": {
230            "array" : T1_err_list,
231            "visible" : True
232         },
233         "xaxis": "x",
234         "yaxis": "y",
235         "mode" : "lines+markers",
236         "line": {"color": color_list[j]},
237         "marker": {"color": color_list[j]},
238         "name" : items,
239         "legendgroup" : str(j),
240         "showlegend" : False
241       }
242            T1_block.append(T1_data)
243
244            T2_data ={
245         "x" : list(exp_dict[items]["residue"]),
246         "y" : T2_list,
247         "error_y": {
248            "array" : T2_err_list,
249            "visible" : True
250         },
251         "xaxis": "x3",
252         "yaxis": "y3",
253         "mode" : "lines+markers",
254         "line": {"color": color_list[j]},
255         "marker": {"color": color_list[j]},
256         "name" : items,
257         "legendgroup": str(j),
258         "showlegend" : False
259       }
260            T2_block.append(T2_data)
261
262            NOE_data ={
263         "x" : list(exp_dict[items]["residue"]),
264         "y" : NOE_list,
265         "error_y": {
266            "array" : NOE_err_list,
267            "visible" : True
268         },
269         "xaxis": "x2",
270         "yaxis": "y2",
271         "mode" : "lines+markers",
272         "line": {"color": color_list[j]},
273         "marker": {"color": color_list[j]},
274         "name" : items,
275         "legendgroup" : str(j),
276         "showlegend" : False
277       }
278            NOE_block.append(NOE_data)
279            j += 1
280
281#flat list of lists
282        list_of_list = [rho_block,T1_block,T2_block,NOE_block]
283        list_of_exp = [item for sublist in list_of_list for item in sublist]
284        exp_plotly = {
285   "data" : list_of_exp,
286   "layout" : {
287       "title" : "Experiment Data Plot",
288       "xaxis" : {
289         "title" : "residue number",
290         "domain" : [0,0.45]
291       },
292       "yaxis" : {
293         "title" : "R1_exp(1/s)",
294         "domain" : [0.55,1]
295       },
296       "xaxis2" : {
297         "title" : "residue number",
298         "domain": [0.55,1],
299         "anchor": "y2"
300       },
301       "yaxis2" : {
302         "title" : "NOE_exp",
303         "domain": [0,0.35],
304         "anchor": "x2"
305       },
306       "xaxis3" : {
307         "title": "residue number",
308         "domain": [0,0.45],
309         "anchor": "y3"
310       },
311       "yaxis3" : {
312         "title": "R2_exp(1/s)",
313         "domain": [0,0.35],
314         "anchor": "x3"
315       },
316       "xaxis4" : {
317         "title": "residue number",
318         "domain": [0.55,1],
319         "anchor": "y4"
320       },
321       "yaxis4" : {
322         "title": "Rho_exp",
323         "domain": [0.65,1],
324         "anchor": "x4"
325       }
326   }
327}
328             
329    else:
330        sys.stderr.write("More than 2 chains are not supported")
331
332    return exp_plotly,exp_keys
Note: See TracBrowser for help on using the repository browser.