/
pymolwrapper.py
87 lines (67 loc) · 3.73 KB
/
pymolwrapper.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
from . import pymol
from .pymol import cmd, cgo
import numpy as np
def start():
pymol.finish_launching()
def draw_spheres(coords, model, frame, radius=0.5):
spheres=[]
for x in coords.reshape(coords.size/3,3):
spheres.extend([cgo.COLOR, 1.0, 0.0, 0.0])
spheres.extend([cgo.SPHERE, x[0], x[1], x[2], radius])
cmd.load_cgo(spheres, model, frame)
def draw_rigid(coords, model, frame, colour, bondslist=[], radius=0.5):
"""
Use pymol to draw a system of rigid body fragments
Parameters
----------
colour: 3-tuple
RBG colour for the spheres being drawn.
bondslist: list of 2-tuples, optional
List of atom pairs between which bonds should be drawn.
"""
spheres = []
for x in coords.reshape(coords.size/3,3):
spheres.extend([cgo.COLOR, colour[0], colour[1], colour[2]])
spheres.extend([cgo.SPHERE, x[0], x[1], x[2], radius])
coords = coords.reshape(coords.size/3,3)
Rcyl = .1
for i in bondslist:
spheres.extend([cgo.CYLINDER, coords[i[0]][0], coords[i[0]][1], coords[i[0]][2],
coords[i[1]][0], coords[i[1]][1], coords[i[1]][2],
Rcyl , 255., 255., 255. , 0., 0., 0.])
cmd.load_cgo(spheres, model, frame)
def draw_box(boxvec, model, frame):
"""
Draw a box around the system to easily visualise periodic boundaries
Parameters
----------
boxvec: np.array
The dimensions of the periodic box
"""
box = []
Rcyl = .1
box.extend([cgo.CYLINDER, -0.5*boxvec[0], -0.5*boxvec[1], -0.5*boxvec[2], 0.5*boxvec[0], -0.5*boxvec[1],
-0.5*boxvec[2], Rcyl, 1., 0., 0., 0., 0., 0.])
box.extend([cgo.CYLINDER, -0.5*boxvec[0], -0.5*boxvec[1], -0.5*boxvec[2], -0.5*boxvec[0], 0.5*boxvec[1],
-0.5*boxvec[2], Rcyl , 0., 1., 0. , 0., 0., 0.])
box.extend([cgo.CYLINDER, -0.5*boxvec[0], -0.5*boxvec[1], -0.5*boxvec[2], -0.5*boxvec[0], -0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 0., 0., 1. , 0., 0., 0.])
box.extend([cgo.CYLINDER, 0.5*boxvec[0], -0.5*boxvec[1], -0.5*boxvec[2], 0.5*boxvec[0], 0.5*boxvec[1],
-0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, 0.5*boxvec[0], -0.5*boxvec[1], -0.5*boxvec[2], 0.5*boxvec[0], -0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, -0.5*boxvec[0], 0.5*boxvec[1], -0.5*boxvec[2], 0.5*boxvec[0], 0.5*boxvec[1],
-0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, -0.5*boxvec[0], 0.5*boxvec[1], -0.5*boxvec[2], -0.5*boxvec[0], 0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, -0.5*boxvec[0], -0.5*boxvec[1], 0.5*boxvec[2], 0.5*boxvec[0], -0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, -0.5*boxvec[0], -0.5*boxvec[1], 0.5*boxvec[2], -0.5*boxvec[0], 0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, 0.5*boxvec[0], 0.5*boxvec[1], 0.5*boxvec[2], 0.5*boxvec[0], 0.5*boxvec[1],
-0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, 0.5*boxvec[0], 0.5*boxvec[1], 0.5*boxvec[2], -0.5*boxvec[0], 0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
box.extend([cgo.CYLINDER, 0.5*boxvec[0], 0.5*boxvec[1], 0.5*boxvec[2], 0.5*boxvec[0], -0.5*boxvec[1],
0.5*boxvec[2], Rcyl , 255., 255., 255. , 0., 0., 0.])
cmd.load_cgo(box, model, frame)