-
Notifications
You must be signed in to change notification settings - Fork 0
/
cctbx_py3dmol_wrapper.py
182 lines (160 loc) · 6.01 KB
/
cctbx_py3dmol_wrapper.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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
from mmtbx.model.model import manager as model_manager
from iotbx.map_manager import map_manager
from iotbx.map_model_manager import map_model_manager
from cubetools import *
import py3Dmol
import numpy as np
import io
class CCTBX3dmolWrapper:
""""
A wrapper around the py3Dmol package written to
easily show CCTBX map/model objects.
"""
def __init__(self,volume_cache_max=6):
"""
3D volume are converted into a string before being send to
the 3dmol javascript library. This is slow, so we cache some
previously calculated volume strings.
"""
self.volume_cache = {}
self.volume_cache_stack = []
self.volume_cache_max = volume_cache_max
def clear_cache(self):
self.volume_cache = {}
self.volume_cache_stack = []
def __call__(self,
*objects,
show_sidechains=True,
show_ribbon=True,
style=None,
ribbon_color="red",
sidechain_color=f"WhiteCarbon",
model_shift = None,
map_color="#808080",
opacity=0.6,
map_threshold=0.1,
map_apix=None,
map_scale=1.0,
shift_back=False,
debug=False):
"""
The main function to show objects. Objects should be of type:
mmtbx.model.model.manager
iotbx.map_manager.map_manager
iotbx.map_model_manager.map_model_manager
When passing multiple objects at once, they will be displayed in a
single viewport.
Parameters:
-----------
objects: multiple cctbx objects to show
show_sidechains (bool,True): Whether to show the protein sidechains
show_ribbon (bool,True): Whether to show the protein ribbon
style (dict,None): The py3Dmol style can be provided explicitly.
ie: {'sphere':{"radius":1.0}} to display all atoms
This will overwrite the other display options.
ribbon_color (str): The color of the ribbons
sidechain_color (str): The color of the sidechains
model_shift (tuple): Modify the model coords by (x,y,z)
map_color (str): The color of the density map
opacity (float): The opacity of the density map
map_threshold (float): The density threshold
map_apix (float): Angstroms per pixel. If None, taken from map_manager
map_scale (float): modify map_apix by a scaler value
shift_back (bool): whether to revert any accumulated shifts
"""
# collect all objects and sort as maps or models
maps = []
models = []
for obj in objects:
if isinstance(obj,map_manager):
maps.append(obj)
elif isinstance(obj, map_model_manager):
models.append(obj.model())
maps.append(obj.map_manager())
elif isinstance(obj,model_manager):
models.append(obj)
else:
print("ERROR: Only high level CCTBX map/model objects are supported by this function")
# objs = [model_manager,map_manager,map_model_manager]
# for obj in objs:
# print("\n",obj.__class__)
return None
view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)
############# Maps #############
map_strings = []
for mm in maps:
n_real = 1
for n in mm.map_data().all():
n_real*=n
if n > 100**3:
print("Map is large:",mm.map_data().all())
print("This may take a very long time or crash")
if map_apix is None:
map_apix = mm.pixel_sizes()
if hash(mm) in self.volume_cache.keys():
cubestring = self.volume_cache[hash(mm)]
if debug:
print("Used volume cache:",True)
else:
data = mm.map_data().as_numpy_array()
meta = self.mm_to_meta(mm)
cubestring = write_cube_string(data,meta)
if debug:
print("Used volume cache:",False)
self.volume_cache[hash(mm)] = cubestring
self.volume_cache_stack.append(hash(mm))
if len(self.volume_cache_stack)>self.volume_cache_max:
del self.volume_cache[self.volume_cache_stack[0]]
self.volume_cache_stack.pop(0)
map_strings.append(cubestring)
for map_string in map_strings:
view.addVolumetricData(map_string, "cube", {'isoval': map_threshold, 'color':map_color , 'opacity': opacity})
############# Models #############
for m in models:
if model_shift is not None:
sites_cart_np = m.get_sites_cart().as_numpy_array()
sites_cart_np+=np.array(model_shift)
m.set_sites_cart(flex.vec3_double(sites_cart_np))
for m in models:
view.addModel(m.model_as_pdb(do_not_shift_back=not shift_back),'pdb')
if style is not None:
view.setStyle(style)
else:
if show_ribbon:
view.setStyle({'cartoon': {'color':ribbon_color}})
else:
BB = ['C','O','N','CA']
view.addStyle({'atom':BB},{'stick':{'colorscheme':sidechain_color,'radius':0.3}})
if show_sidechains:
BB = ['C','O','N']
view.addStyle({'and':[{'resn':["GLY","PRO"],'invert':True},{'atom':BB,'invert':True}]},
{'stick':{'colorscheme':sidechain_color,'radius':0.3}})
view.addStyle({'and':[{'resn':"GLY"},{'atom':'CA'}]},
{'sphere':{'colorscheme':sidechain_color,'radius':0.3}})
view.addStyle({'and':[{'resn':"PRO"},{'atom':['C','O'],'invert':True}]},
{'stick':{'colorscheme':sidechain_color,'radius':0.3}})
view.zoomTo()
return view
@staticmethod
def mm_to_meta(map_manager,scale=1.0):
"""
Form metadata dictionary from map manager object
mm: map_manager objects
returns: meta (dict)
"""
conversion = 1.8897259885789233 # bohr to angstrom
scale*= conversion
ax,ay,az = map_manager.pixel_sizes()
ax*=scale
ay*=scale
az*=scale
xvec = (ax,0.0,0.0)
yvec = (0.0,ay,0.0)
zvec = (0.0,0.0,az)
origin = tuple(map_manager.get_origin())
meta = {"atoms":[],
"org":origin,
"xvec":xvec,
"yvec":yvec,
"zvec":zvec}
return meta