source: rotdif/bin/get_elm_tensor.py @ 1705

Last change on this file since 1705 was 1705, checked in by yuexi, 2 years ago

finish ellipsoid

File size: 2.2 KB
Line 
1import re
2import operator
3
4def elm_tensor(pdb, scale=10):
5    with open('ELM_prediction') as in_f:
6        tmp = in_f.readlines()
7    for i in range(len(tmp)):
8        if tmp[i].startswith('====Diffusion Tensor Sorted Eigendecomposition===='):
9            row_dx = [float(it)*scale for it in re.findall(r"[-+]?\d*\.\d+|[-+]?\d+",tmp[i+1])[:3]]
10            row_dy = [float(it)*scale for it in re.findall(r"[-+]?\d*\.\d+|[-+]?\d+",tmp[i+2])[:3]]
11            row_dz = [float(it)*scale for it in re.findall(r"[-+]?\d*\.\d+|[-+]?\d+",tmp[i+3])[:3]]
12    dx = [row_dx[0],row_dy[0],row_dz[0]]
13    dy = [row_dx[1],row_dy[1],row_dz[1]]
14    dz = [row_dx[2],row_dy[2],row_dz[2]]
15    with open(pdb) as in_f2:
16        tmp2 = in_f2.readlines()
17    ori_x = 0
18    ori_y = 0
19    ori_z = 0
20    count = 0
21    for j in range(len(tmp2)):
22        if tmp2[j].startswith('ATOM'):
23            count += 1
24            tmp_str = ''.join(tmp2[j])
25            tmp_str = " ".join(tmp_str.split())
26            tmp_ls = tmp_str.split(" ")
27            ori_x += float(tmp_ls[6])
28            ori_y += float(tmp_ls[7])
29            ori_z += float(tmp_ls[8])
30   
31    ori_x = ori_x / count
32    ori_y = ori_y / count
33    ori_z = ori_z / count
34    ori_coor = [ori_x, ori_y, ori_z]
35    dx_1 = list(map(operator.sub, ori_coor , dx ))
36    dx_2 = list(map(operator.add, ori_coor , dx ))
37    dy_1 = list(map(operator.sub, ori_coor , dy ))
38    dy_2 = list(map(operator.add, ori_coor , dy ))
39    dz_1 = list(map(operator.sub, ori_coor , dz ))
40    dz_2 = list(map(operator.add, ori_coor , dz ))
41    all_x = [str(i) for i in dx_1+dx_2]
42    all_y = [str(i) for i in dy_1+dy_2]
43    all_z = [str(i) for i in dz_1+dz_2]
44    obj = "[ CYLINDER, %s, 0.2, 1.0, 1.0, 1.0, 1.0, 0.0, 0.,\
45         CYLINDER, %s, 0.2, 1.0, 1.0, 1.0, 0., 1.0, 0.,\
46        CYLINDER, %s,0.2, 1.0, 1.0, 1.0, 0., 0.0, 1.0,]" % (','.join(all_x), ','.join(all_y), ','.join(all_z))
47   
48    with open('ELM_tensor_axes.py','w') as out_f:
49        out_f.write("""# axes plot script
50from pymol.cgo import *
51from pymol import cmd
52from pymol.vfont import plain
53#create the axes object, draw axes with cylinders colored red (Dx), green (Dy), blue (Dz) \n""" +
54"obj = " + obj +
55"""\n # then we load it into PyMOL \n
56cmd.load_cgo(obj,'out_ani')""")   
Note: See TracBrowser for help on using the repository browser.