forked from gnina/scripts
-
Notifications
You must be signed in to change notification settings - Fork 0
/
show_xyz_arrows.py
executable file
·103 lines (75 loc) · 2.63 KB
/
show_xyz_arrows.py
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
#!/usr/bin/env python3
'''Adds show_arrows command to pymol, which takes an xyz file'''
import sys
import os
from pymol import cmd, cgo, CmdException
from chempy import cpv
def draw_arrow(xyz1,xyz2, radius=0.5, gap=0.0, hlength=-1, hradius=-1,
color='blue red', name=''):
'''
Draw an arrow; borrows heavily from cgi arrows.
'''
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
xyz1 = list(xyz1)
xyz2 = list(xyz2)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def make_pymol_arrows(base, atoms, scale, color, radius):
arrow_objs = []
arrow_group = base + '_arrows'
cmd.delete(arrow_group) #remove any pre-existing group
for i, atom in enumerate(atoms):
arrow_obj = base + '_arrow_' + str(i)
arrow_objs.append(arrow_obj)
elem, xi, yi, zi, dx, dy, dz = atom
c = 1.725*radius
xf = xi + -scale*dx + c
yf = yi + -scale*dy + c
zf = zi + -scale*dz + c
draw_arrow((xi,yi,zi),(xf,yf,zf),radius=radius,color=color,name=arrow_obj)
cmd.group(arrow_group,' '.join(arrow_objs))
def xyz_line_to_atom(xyz_line):
fields = xyz_line.split()
elem = fields[0]
x = float(fields[1])
y = float(fields[2])
z = float(fields[3])
dx = float(fields[4])
dy = float(fields[5])
dz = float(fields[6])
return elem, x, y, z, dx, dy, dz
def read_xyz_file(xyz_file):
with open(xyz_file, 'r') as f:
lines = f.readlines()
n_atoms = int(lines[0])
atoms = []
for i in range(n_atoms):
atom = xyz_line_to_atom(lines[2+i])
atoms.append(atom)
return atoms
def show_xyz_arrows(xyzfile, scale=2.0, color="white purple",radius=0.2):
atoms = read_xyz_file(xyzfile)
base_name = xyzfile.replace('.xyz', '')
make_pymol_arrows(base_name, atoms, float(scale), color, float(radius))
cmd.extend('show_xyz_arrows', show_xyz_arrows)