/
cgo_arrow.py
82 lines (56 loc) · 2.01 KB
/
cgo_arrow.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
'''
http://pymolwiki.org/index.php/cgo_arrow
(c) 2013 Thomas Holder, Schrodinger Inc.
License: BSD-2-Clause
'''
from pymol import cmd, cgo, CmdException
def cgo_arrow(atom1='pk1', atom2='pk2', radius=0.5, gap=0.0, hlength=-1, hradius=-1,
color='black black', name=''):
'''
#I modify the color the line just before
DESCRIPTION
Create a CGO arrow between two picked atoms.
ARGUMENTS
atom1 = string: single atom selection or list of 3 floats {default: pk1}
atom2 = string: single atom selection or list of 3 floats {default: pk2}
radius = float: arrow radius {default: 0.5}
gap = float: gap between arrow tips and the two atoms {default: 0.0}
hlength = float: length of head
hradius = float: radius of head
color = string: one or two color names {default: blue red}
name = string: name of CGO object
'''
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
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)
cmd.extend('cgo_arrow', cgo_arrow)