def view3d(i=0):
'''Simple structure viewer that uses py3Dmol to view PDB structure by
indexing the list of PDBids
Parameters
----------
i (int): index of the protein if a list of PDBids
'''
print(f"PdbID: {pdbIds[i]}, Style: {style}")
if '.' not in pdbIds[i]:
viewer = py3Dmol.view(query='pdb:' + pdbIds[i],
options={'doAssembly': bioAssembly})
viewer.setStyle({style: {'color': color}})
viewer.setStyle({'hetflag': True},
{'stick': {
'singleBond': False
}})
else:
pdbid, chainid = pdbIds[i].split('.')
viewer = py3Dmol.view(query='pdb:' + pdbid,
options={'doAssembly': bioAssembly})
viewer.setStyle({})
viewer.setStyle({'chain': chainid}, {style: {'color': color}})
viewer.setStyle({
'chain': chainid,
'hetflag': True
}, {'stick': {
'singleBond': False
}})
viewer.zoomTo({'chain': chainid})
return viewer.show()
def show_sticks(xyz=None, species=None, project_directory=None):
"""
Draws the molecule in a "sticks" style according to the supplied xyz coordinates
Returns whether successful of not
If successful, save an image using draw_3d
"""
xyz = check_xyz_species_for_drawing(xyz, species)
coordinates, _, _, _, _ = get_xyz_matrix(xyz)
s_mol, b_mol = molecules_from_xyz(xyz)
mol = b_mol if b_mol is not None else s_mol
try:
_, rd_mol, _ = rdkit_conf_from_mol(mol, coordinates)
except ValueError:
return False
mb = Chem.MolToMolBlock(rd_mol)
p = p3d.view(width=400, height=400)
p.addModel(mb, 'sdf')
p.setStyle({'stick': {}})
# p.setBackgroundColor('0xeeeeee')
p.zoomTo()
p.show()
draw_3d(xyz=xyz,
species=species,
project_directory=project_directory,
save_only=True)
return True
def view3d(i=0):
'''Simple structure viewer that uses py3Dmol to view PDB structure by
indexing the list of PDBids
Attributes
----------
i (int): index of the protein if a list of PDBids
'''
structures = df['pdbId'].iloc
groups = df['groupNum0'].iloc
chains = df['chain0'].iloc
elements = df['element0'].iloc
ori = str(df[sortBy].iloc[i])[:5]
print(f"PDBId: {structures[i]} chain: {chains[i]} element: {elements[i]}")
print(f"{sortBy}: {ori}")
viewer = py3Dmol.view(query='pdb:' + structures[i], width=700, height=700)
neighbors = {'resi': get_neighbors_group(i), 'chain': get_neighbors_chain(i)}
metal = {'resi': groups[i], 'atom': str(elements[i]).upper(), 'chain': chains[i]}
viewer.setStyle(neighbors, {'stick': {'colorscheme': 'orangeCarbon'}})
viewer.setStyle(metal, {'sphere': {'radius': 0.5, 'color': 'gray'}})
viewer.zoomTo(neighbors)
return viewer.show()
def to_3Dmol(self, style=None, **kwargs):
"""Generate protein structure & return interactive py3Dmol.view for visualization.
Args:
style (str, optional): Style string to be passed to py3Dmol for
visualization. Defaults to None.
Returns:
py3Dmol.view object: A view object that is interactive in iPython notebook
settings.
"""
import py3Dmol
if not style:
style = {
'cartoon': {
'color': 'spectrum'
},
'stick': {
'radius': .15
}
}
self._initialize_coordinates_and_PdbCreator()
view = py3Dmol.view(**kwargs)
view.addModel(self.pdb_creator.get_pdb_string(), 'pdb')
if style:
view.setStyle(style)
view.zoomTo()
return view
def show_structure(domain, full=False):
'''
Uses py3Dmol and requests to fetch domain PDB and show it in jupyter
'''
view = py3Dmol.view()
def adjust_looks(view):
view.setStyle({'cartoon': {'colorscheme': 'ssJmol'}})
view.center()
return view
if full:
prot = domain[:4].upper()
if (prot + '.pdb') in os.listdir('structures'):
with open('./structures/' + prot + '.pdb') as file:
view.addModel(file.read(), 'pdb')
else:
r = requests.get('https://files.rcsb.org/download/' + prot +
'.pdb')
with open('./structures/' + prot + '.pdb', 'w') as file:
file.write(r.text)
view = view.addModel(r.text, 'pdb')
else:
if (domain + '.pdb') in os.listdir('structures'):
with open('./structures/' + domain + '.pdb') as file:
view.addModel(file.read(), 'pdb')
else:
r = requests.get(
'http://www.cathdb.info/version/v4_2_0/api/rest/id/' + domain +
'.pdb')
with open('./structures/' + domain + '.pdb', 'w') as file:
file.write(r.text)
view = view.addModel(r.text, 'pdb')
return adjust_looks(view)
def display_molecules(molecules, size=(400, 300), animation=False):
"""
display the molecules in Pymatgen object.
Args:
molecules: a list of pymatgen molecules
size: (width, height) in tuple
animation: whether or not display animation
Returns:
py3Dmol object
"""
(width, height) = size
view = py3Dmol.view(height=height, width=width)
mol_strs = ""
for mol in molecules:
mol_strs += mol.to(fmt='xyz') + '\n'
if animation:
view.addModelsasFrames(mol_strs, 'xyz')
view.animate({'loop': 'forward'})
else:
view.addModels(mol_strs, 'xyz')
view.setStyle({'stick': {'colorscheme': 'greenCarbon'}})
return view.zoomTo()
def view3d(i=0):
'''
Simple structure viewer that zooms into a specified group and highlight
its neighbors
Attributes:
i (int): index of the protein if a list of PDBids
'''
print(
f"PDB: {pdbIds[i]}, group: {groups[i]}, chain: {chains[i]}, cutoffDistance: {distance}"
)
center = {'resi': groups[i], 'chain': chains[i]}
neighbors = {
'resi': groups[i],
'chain': chains[i],
'byres': 'true',
'expand': distance
}
viewer = py3Dmol.view(query='pdb:' + pdbIds[i])
viewer.zoomTo(center)
viewer.setStyle(neighbors, {'stick': {}})
viewer.setStyle(center, {'sphere': {'color': 'red'}})
viewer.zoom(0.2, 1000)
return viewer.show()
def drawConformer_episodic(data: dict,
confIds: List[int],
size: Tuple[int, int] = (300, 300),
style: str = "stick") -> py3Dmol.view:
"""Displays a specified conformer for each episode loaded in `data`.
Parameters
----------
data : dict from string to list
Contains the loaded episode information. 'mol' must be a key in data and the corresponding list must contain
RDKit Mol objects.
confIds : list of int
The indices for the conformers to be displayed (for each episode loaded in data).
size : Tuple[int, int]
The size of the display for each individual molecule (width, height).
style: str
The drawing style for displaying the molecule. Can be sphere, stick, line, cross, cartoon, and surface.
"""
n = len(data['mol'])
view = py3Dmol.view(width=size[0] * n,
height=size[0] * n,
linked=False,
viewergrid=(n, 1))
for i in range(n):
block = Chem.MolToMolBlock(data['mol'][i], confId=confIds[i])
view.addModel(block, 'mol', viewer=(i, 0))
view.setStyle({style: {}}, viewer=(i, 0))
view.zoomTo()
return view
def __init__(self, xsize, ysize, nrow=1, ncol=1):
from py3Dmol import view
# set linked to False so that each structure can be rotated
# independently
self.xyzview = view(width=ncol * xsize,
height=nrow * ysize,
viewergrid=(nrow, ncol),
linked=False)
def display_pdb(path_to_pdb):
with open(path_to_pdb) as ifile:
protein = "".join([x for x in ifile])
view = py3Dmol.view(width=400, height=300)
view.addModelsAsFrames(protein)
view.setStyle({'model': -1}, {"cartoon": {'color': 'spectrum'}})
view.zoomTo()
view.show()
def draw3d(mol_sd):
fig = py3Dmol.view(width=600, height=400)
if isinstance(mol_sd, str):
fig.addModel(mol_sd, 'sdf')
elif isinstance(mol_sd, list):
for m in mol_sd:
fig.addModel(m, 'sdf')
fig.setStyle({'stick': {'colorscheme': 'grayCarbon'}})
return fig.render()
def drawit(m,dimensions=(900,400),p=None,confId=-1):
mb = Chem.MolToMolBlock(m,confId=confId)
if p is None:
p = py3Dmol.view(width=dimensions[0],height=dimensions[1])
p.removeAllModels()
p.addModel(mb,'sdf')
p.setStyle({'stick':{}})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
return p.show()
def display_atomic(struc,
size=(600, 300),
scale=0.25,
radius=0.10,
supercell=(1, 1, 1),
show_wp=False):
"""
display the molecular crystals generated from pyxtal. If the animation is False,
only dump the structure to cif and let py3Dmol display it. If animation is on,
show the generation of molecules in the crystal as an animation.
Args:
size: (width, height) in tuple
scale: the size of sphere
radius: the size of stick
supercell: replicate the crystal (valid only when animation is False)
show_wp: whether or not highlight the unique wyckoff site
Returns:
py3Dmol object
"""
(width, height) = size
view = py3Dmol.view(height=height, width=width)
if struc.dim == 0:
fmt = "xyz"
else:
fmt = "cif"
txt = struc.to_file(fmt=fmt)
view.addModel(txt, fmt, {
'doAssembly': True,
'duplicateAssemblyAtoms': True
})
view.setStyle({
'sphere': {
'colorscheme': 'Jmol',
'scale': scale
},
'stick': {
'colorscheme': 'Jmol',
'radius': radius
}
})
if struc.dim != 0:
view.addUnitCell()
A, B, C = supercell
view.replicateUnitCell(A, B, C)
if show_wp:
view.setStyle({'sym': 2},
{'sphere': {
'scale': scale * 1.1,
'color': 'blue'
}})
return view.zoomTo()
def _view_3Dmol(self, **options):
"""
Visualize the system using 3Dmol.js and py3Dmol.
"""
from nappy.elements import elements
try:
import py3Dmol
except Exception as e:
raise
if 'model' in options.keys():
modopts = options['model']
if 'volume' in options.keys():
volopts = options['volume']
pdb = nappy.io.get_PDB_txt(self, **options)
v = py3Dmol.view()
v.removeAllModels()
v.setViewStyle({'style': 'outline', 'color': 'black', 'width': 0.03})
v.addModel(pdb, 'pdb')
colors = {s: elements[s]['CPKcolor'] for s in self.specorder}
if 'modopts' in locals():
v.setStyle(modopts)
else:
v.setStyle({
'stick': {
'radius': .1,
'colorscheme': {
'prop': 'elem',
'map': colors
}
},
'sphere': {
'radius': .5,
'colorscheme': {
'prop': 'elem',
'map': colors
}
}
})
v.addUnitCell()
if hasattr(self, 'voldata'):
cube = self.get_cube_txt()
if 'volopts' in locals():
v.addVolumetricData(cube, "cube", volopts)
else:
v.addVolumetricData(cube, "cube", {
'isoval': 300,
'color': "red",
'opacity': 0.95
})
v.zoomTo()
v.show()
return None
def view(self):
"""
A method designed to create a 3D figure of the AutoTST_Molecule with py3Dmol from the rdkit_molecule
"""
mb = Chem.MolToMolBlock(self.rdkit_molecule)
p = py3Dmol.view(width=600, height=600)
p.addModel(mb, "sdf")
p.setStyle({'stick': {}})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
return p.show()
def vis_xyz(file_path):
import py3Dmol
with open(file_path) as f:
xyz = f.read()
xyzview = py3Dmol.view(width=400, height=400)
xyzview.addModel(xyz, 'xyz')
xyzview.setStyle({'stick': {}})
xyzview.setBackgroundColor('0xeeeeee')
xyzview.zoomTo()
xyzview.show()
def _graph_conf(m, confId=0, energies=[]):
mb = Chem.MolToMolBlock(m, confId=confId)
p = py3Dmol.view(width=400, height=400)
p.removeAllModels()
p.addModel(mb, 'sdf')
p.setStyle({'stick': {}})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
if len(energies) > 0:
print(
f"G: {min(energies):.2f} + {energies[confId] - min(energies):.2f} kcal/mol"
)
return p.show()
def display_molecular(struc,
size=(600, 300),
supercell=(1, 1, 1),
axis=None,
animation=False,
interval=2000):
"""
display the molecular crystals generated from pyxtal. If the animation is False,
only dump the structure to cif and let py3Dmol display it. If animation is on,
show the generation of molecules in the crystal as an animation.
Args:
size: (width, height) in tuple
supercell: replicate the crystal (valid only when animation is False)
axis: 3-vector to reprent the rotational axis
animation: whether or not display the animation
Returns:
py3Dmol object
"""
(width, height) = size
view = py3Dmol.view(height=height, width=width)
if not animation:
cif = struc.to_file()
view.addModel(
cif, 'cif', {
'doAssembly': True,
'duplicateAssemblyAtoms': True,
'normalizeAssembly': True
})
view.setStyle({'stick': {'colorscheme': 'greenCarbon'}})
view.addUnitCell()
if axis is not None:
site = struc.mol_sites[0]
for id in range(len(site.wp)):
mol = site.get_mol_object(id)
addlines(view, mol.cart_coords.mean(axis=0), [axis.copy()])
A, B, C = supercell
view.replicateUnitCell(A, B, C)
else:
cifs = ""
for i in range(len(struc.mol_sites[0].wp)):
cifs += struc.to_file(sym_num=i + 1)
view.addModelsAsFrames(cifs, 'cif')
view.animate({'loop': 'forward', 'interval': interval})
view.addUnitCell()
view.setStyle({'model': 0}, {'stick': {'colorscheme': 'greenCarbon'}})
return view.zoomTo()
def draw_3d(self):
try:
import py3Dmol
AllChem.EmbedMultipleConfs(self.rdkitmol_Hs)
mb = Chem.MolToMolBlock(self.rdkitmol_Hs)
p = py3Dmol.view(width=300,height=300)
p.addModel(mb,'sdf')
p.setStyle({'stick':{}})
# Styles: stick, line, cross, sphere
p.zoomTo()
p.show()
return p
except:
return 'py3Dmol and rdkit required'
def view_ts(self, mol=None):
"""
A method designed to create a 3D figure of the Multi_Molecule with py3Dmol
"""
if not mol:
mol = self.rdkit_ts
mb = Chem.MolToMolBlock(mol)
p = py3Dmol.view(width=400, height=400)
p.addModel(mb, "sdf")
p.setStyle({'stick': {}})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
return p.show()
def view3d(i=0):
'''
Simple structure viewer that uses py3Dmol to view PDB structure by
indexing the list of PDBids
Attributes:
i (int): index of the protein if a list of PDBids
'''
print(f"PdbID: {pdbIds[i]}, Style: {style}")
viewer = py3Dmol.view(query='pdb:' + pdbIds[i])
viewer.setStyle({style: {'color': color}})
return viewer.show()
def draw_all(rdmol):
def draw_one(m, p, confId=-1):
mb = Chem.MolToMolBlock(m, confId=confId)
p.removeAllModels()
p.addModel(mb, 'sdf')
p.setStyle({'stick': {}})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
return p.show()
p = py3Dmol.view(width=400, height=400)
interact(draw_one,
m=fixed(rdmol),
p=fixed(p),
confId=(0, rdmol.GetNumConformers() - 1))
def draw_mols_surfs(mols,
width=400,
height=400,
surface=True,
surface_opacity=0.5):
view = py3Dmol.view(width=width, height=height)
view.setBackgroundColor('0xeeeeee')
view.removeAllModels()
for mol in mols:
addMolToView(mol, view)
if surface:
view.addSurface(py3Dmol.SAS, {'opacity': surface_opacity})
view.zoomTo()
return view.show()
def jupyter_draw_conf(m, p=None, confId=-1, style='stick', extrastyle={}):
"""
remember to
from rdkit.Chem.Draw import IPythonConsole
"""
mb = Chem.MolToMolBlock(m, confId=confId)
if p is None:
p = py3Dmol.view(width=400, height=400)
p.removeAllModels()
p.addModel(mb, 'sdf')
p.setStyle({style: extrastyle})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
return p.show()
def show_xyz(self, width=600, height=600, print_xyz=False):
""" Displays a 3D rendering of the conformer using Py3Dmol
:param width: the width of the display window
:param height: the height of the display window
"""
XYZ = "{0:3d}\n{1:s}\n".format(self.NAtoms, self._id)
for at, xyz in zip(self.atoms, self.xyz):
XYZ += "{0:3s}{1:10.3f}{2:10.3f}{3:10.3f}\n".format(
at, xyz[0], xyz[1], xyz[2])
xyzview = p3D.view(width=width, height=height)
xyzview.addModel(XYZ, 'xyz')
xyzview.setStyle({'stick': {}})
xyzview.zoomTo()
xyzview.show()
if print_xyz == True: print(XYZ)
def show(self, **kwargs):
"""
Display 3Dmol.js output using IPython/Jupyter notebook display.
If called with keyword arguments, those will be passed on to
`update()`.
"""
self.view = py3Dmol.view(width=self.width, height=self.height)
for i, m in enumerate(self.models):
m.add_to_view(self.view, i)
self.view.zoomTo()
self._set_hoverable()
self.view.show()
if len(kwargs) > 0:
self.update(**kwargs)
def plot(smiles, size=(200, 200), style='stick'):
"""
Visualize a molecule from its formula (smiles)
Adapted from https://birdlet.github.io/2019/10/02/py3dmol_example
"""
mol = Chem.MolFromSmiles(smiles)
mol = Chem.AddHs(mol)
AllChem.EmbedMolecule(mol)
AllChem.MMFFOptimizeMolecule(mol)
block = Chem.MolToMolBlock(mol)
viewer = py3Dmol.view(width=size[0], height=size[1])
viewer.addModel(block, 'mol')
viewer.setStyle({style: {}})
viewer.zoomTo()
viewer.show()
def draw(mol: Chem.Mol, energies: list = []) -> interact:
"""Make a drawing of all conformers in 3d
:param mol: rdkit molecule
:type mol: rdkit.Chem.Mol
:param energies: list of conformer energies
:type energies: list
:return: interact 3D object
"""
p = py3Dmol.view(width=400, height=400)
return interact(_graph_conf,
m=fixed(mol),
p=fixed(p),
confId=[c.GetId() for c in mol.GetConformers()],
energies=fixed(energies))
def show_3D_mol(mol, confId):
"""
Shows a molecule in 3D
:param: an RDKit.mol
:param confId: the IR of the conformer to show
:return: molBlock
"""
mb = Chem.MolToMolBlock(mol, confId=confId)
p = py3Dmol.view()
p.addModel(mb, 'sdf')
p.setStyle({'stick': {}})
p.setBackgroundColor('0xeeeeee')
p.zoomTo()
p.show()
return mb
def drawMol3D(m, view=None, confId=-1, drawAs=None, bgColor=None, size=None):
if bgColor is None:
bgColor = bgcolor_3d
if size is None:
size = molSize_3d
if view is None:
view = py3Dmol.view(width=size[0], height=size[1])
view.removeAllModels()
try:
ms = iter(m)
for m in ms:
addMolToView(m, view, confId, drawAs)
except TypeError:
addMolToView(m, view, confId, drawAs)
view.setBackgroundColor(bgColor)
view.zoomTo()
return view.show()
def show(
self,
*,
style: Union[str, Dict[str, Any]] = "ball_and_stick",
canvas: Tuple[int, int] = (400, 400)
) -> 'py3Dmol.view':
"""Creates a 3D representation of a moleucle that can be manipulated in Jupyter Notebooks and exported as images (`.png`).
Parameters
----------
style : Union[str, Dict[str, Any]], optional
Either 'ball_and_stick' or 'stick' style representations or a valid py3Dmol style dictionary.
canvas : Tuple[int, int], optional
The (width, height) of the display canvas in pixels
Returns
-------
py3Dmol.view
A py3dMol view object of the molecule
"""
if not which_import("py3Dmol", return_bool=True):
raise ModuleNotFoundError(
f"Python module py3DMol not found. Solve by installing it: `conda install -c conda-forge py3dmol`"
) # pragma: no cover
import py3Dmol
if isinstance(style, dict):
pass
elif style == 'ball_and_stick':
style = {'stick': {'radius': 0.2}, 'sphere': {'scale': 0.3}}
elif style == 'stick':
style = {'stick': {}}
else:
raise KeyError(f"Style '{style}' not recognized.")
xyzview = py3Dmol.view(width=canvas[0], height=canvas[1])
xyzview.addModel(self.to_string("xyz", units="angstrom"), 'xyz')
xyzview.setStyle(style)
xyzview.zoomTo()
return xyzview
def drawMol3D(m, view=None, confId=-1, drawAs=None, bgColor=None, size=None):
if bgColor is None:
bgColor = bgcolor_3d
if size is None:
size = molSize_3d
if view is None:
view = py3Dmol.view(width=size[0], height=size[1])
view.removeAllModels()
try:
iter(m)
except TypeError:
addMolToView(m, view, confId, drawAs)
else:
ms = m
for m in ms:
addMolToView(m, view, confId, drawAs)
view.setBackgroundColor(bgColor)
view.zoomTo()
return view.show()
def view3D(*alist, **kwargs):
"""Return a py3Dmol view instance for interactive visualization in
Jupyter notebooks. Available arguments are: width, height (of
the viewer), backgroundColor, zoomTo (a py3Dmol selection to center
around), and style, which is a py3Dmol style object that will be
applied to all atoms in the scene. More complex styling can be achieved
by manipulating the view object directly.
The default style is to show the protein in a rainbow cartoon and
hetero atoms in sticks/spheres.
GNM/ANM Coloring
An array of fluctuation values (flucts) can be provided with the flucts kwarg
for visualization of GNM/ANM calculations. The array is assumed to
correpond to a calpha selection of the provided protein.
The default color will be set to a RWB color scheme on a per-residue
basis. If the fluctuation vector contains negative values, the
midpoint (white) will be at zero. Otherwise the midpoint is the mean.
An array of displacement vectors (vecs) can be provided with the vecs kwarg.
The animation of these motions can be controlled with frames (number
of frames to animate over), amplitude (scaling factor), and animate
(3Dmol.js animate options).
If multiple structures are provided with the flucts or vecs arguments, these
arguments must be provided as lists of arrays of the appropriate dimension.
"""
try:
import py3Dmol
except:
raise ImportError('py3Dmol needs to be installed to use view3D')
from .pdbfile import writePDBStream
width = kwargs.get('width', 400)
height = kwargs.get('height', 400)
data_list = kwargs.pop('data', None)
modes = kwargs.pop('mode', None)
style = kwargs.pop('style', [])
zoomto = kwargs.pop('zoomto', {})
bgcolor = kwargs.pop('backgroundcolor', 'white')
bgcolor = kwargs.pop('backgroundColor', bgcolor)
frames = kwargs.pop('frames', 30)
interval = kwargs.pop('interval', 1)
anim = kwargs.pop('anim', False)
scale = kwargs.pop('scale', 100)
if modes is None:
n_modes = 0
else:
modes = wrapModes(modes)
n_modes = len(modes)
if data_list is None:
n_data = 0
else:
data_list = wrapModes(data_list)
n_data = len(data_list)
view = py3Dmol.view(width=width, height=height, js=kwargs.get('js','http://3dmol.csb.pitt.edu/build/3Dmol-min.js'))
def _mapData(atoms, data):
# construct map from residue to data property
propmap = []
for j, a in enumerate(atoms.calpha):
propmap.append({'model': -1, 'chain': a.getChid(), 'resi': int(a.getResnum()),
'props': {'data': float(data[j]) } })
# set the atom property
# TODO: implement something more efficient on the 3Dmol.js side (this is O(n*m)!)
view.mapAtomProperties(propmap)
# color by property using gradient
extreme = np.max(np.abs(data))
lo = -extreme if np.min(data) < 0 else 0
mid = np.mean(data) if np.min(data) >= 0 else 0
view.setColorByProperty({'model': -1}, 'data', 'rwb', [extreme,lo,mid])
view.setStyle({'model': -1},{'cartoon':{'style':'trace'}})
for i, atoms in enumerate(alist):
pdb = createStringIO()
writePDBStream(pdb, atoms)
view.addAsOneMolecule(pdb.getvalue(), 'pdb')
view.setStyle({'model': -1}, {'cartoon': {'color':'spectrum'}})
view.setStyle({'model': -1, 'hetflag': True}, {'stick':{}})
view.setStyle({'model': -1, 'bonds': 0}, {'sphere':{'radius': 0.5}})
if n_data:
data = data_list[i]
try:
data = data.getArray()
except AttributeError:
pass
if atoms.calpha.numAtoms() != len(data):
raise RuntimeError("Atom count mismatch: {} vs {}. data styling assumes a calpha selection."
.format(atoms.calpha.numAtoms(), len(data)))
_mapData(atoms, data)
if n_modes:
mode = modes[i]
try:
arr = mode.getArray()
is3d = mode.is3d()
except AttributeError:
arr = mode
if is3d:
if atoms.calpha.numAtoms()*3 != len(arr):
raise RuntimeError("Atom count mismatch: {} vs {}. mode animation assumes a calpha selection."
.format(atoms.calpha.numAtoms(), len(arr)//3))
if anim:
# construct map from residue to anm property and dx,dy,dz vectors
propmap = []
for j, a in enumerate(atoms.calpha):
propmap.append({'model': -1, 'chain': a.getChid(), 'resi': int(a.getResnum()),
'props': {'dx': arr[3*j], 'dy': arr[3*j+1], 'dz': arr[3*j+2] } })
# set the atom property
# TODO: implement something more efficient on the 3Dmol.js side (this is O(n*m)!)
view.mapAtomProperties(propmap)
else:
for j, a in enumerate(atoms.calpha):
start = a._getCoords()
dcoords = arr[3*j:3*j+3]
end = start + dcoords * scale
view.addArrow({'start': {'x':start[0], 'y':start[1], 'z':start[2]},
'end': {'x':end[0], 'y':end[1], 'z':end[2]},
'radius': 0.3})
else:
if atoms.calpha.numAtoms() != len(arr):
raise RuntimeError("Atom count mismatch: {} vs {}. mode styling assumes a calpha selection."
.format(atoms.calpha.numAtoms(), len(arr)))
_mapData(atoms, arr)
# setting styles ...
view.setBackgroundColor(bgcolor)
if n_modes:
#create vibrations
view.vibrate(frames, scale)
animate = kwargs.get('animate', {'loop':'rock', 'interval':interval})
view.animate(animate)
if isinstance(style, dict):
style = ({}, style)
if isinstance(style, tuple):
styles = [style]
else:
styles = style
for sel, style in styles:
view.setStyle(sel, style)
view.zoomTo(zoomto)
return view
def view3D(*alist, **kwargs):
"""Return a py3Dmol view instance for interactive visualization in
Jupyter notebooks. Available arguments are: width, height (of
the viewer), backgroundColor, zoomTo (a py3Dmol selection to center
around), and style, which is a py3Dmol style object that will be
applied to all atoms in the scene. More complex styling can be achieved
by manipulating the view object directly.
The default style is to show the protein in a rainbow cartoon and
hetero atoms in sticks/spheres.
GNM/ANM Coloring
An array of fluctuation values can be provided with the flucts kwarg
for visualization of GNM/ANM calculations. The array is assumed to
correpond to a calpha selection of the provided protein.
The default color will be set to a RWB color scheme on a per-residue
basis. If the fluctuation vector contains negative values, the
midpoint (white) will be at zero. Otherwise the midpoint is the mean.
An array of displacement vectors can be provided with the vecs kwarg.
The animation of these motions can be controlled with frames (number
of frames to animate over), amplitude (scaling factor), and animate
(3Dmol.js animate options).
"""
import StringIO, py3Dmol
from pdbfile import writePDBStream
pdb = StringIO.StringIO()
for atoms in alist:
writePDBStream(pdb, atoms)
width = kwargs.get('width',400)
height = kwargs.get('height',400)
view = py3Dmol.view(width=width,height=height,js=kwargs.get('js','http://3dmol.csb.pitt.edu/build/3Dmol-min.js'))
#case insensitive kwargs..
bgcolor = kwargs['backgroundcolor'] if 'backgroundcolor' in kwargs else kwargs.get('backgroundColor','white')
view.setBackgroundColor(bgcolor)
view.addModels(pdb.getvalue(),'pdb')
view.setStyle({'cartoon': {'color':'spectrum'}})
view.setStyle({'hetflag': True}, {'stick':{}})
view.setStyle({'bonds': 0}, {'sphere':{'radius': 0.5}})
if 'flucts' in kwargs:
garr = kwargs['flucts']
#note we are only getting info from last set of atoms..
if atoms.calpha.numAtoms() != len(garr):
raise RuntimeError("Atom count mismatch: {} vs {}. flucts styling assume a calpha selection.".format(atoms.calpha.numAtoms(), len(garr)))
else:
#construct map from residue to flucts property
propmap = []
for (i,a) in enumerate(atoms.calpha):
propmap.append({'chain': a.getChid(), 'resi':a.getResnum(), 'props': {'flucts': garr[i] } })
#set the atom property
#TODO: implement something more efficient on the 3Dmol.js side (this is O(n*m)!)
view.mapAtomProperties(propmap)
#color by property using gradient
extreme = np.abs(garr).max()
lo = -extreme if garr.min() < 0 else 0
mid = np.mean(garr) if garr.min() >= 0 else 0
view.setColorByProperty({}, 'flucts', 'rwb', [extreme,lo,mid])
view.setStyle({'cartoon':{'style':'trace'}})
if 'vecs' in kwargs:
aarr = kwargs['vecs'] #has xyz coordinates
#note we are only getting info from last set of atoms..
if atoms.calpha.numAtoms()*3 != len(aarr):
raise RuntimeError("Atom count mismatch: {} vs {}. vecs animation assume a calpha selection.".format(atoms.calpha.numAtoms(), len(aarr)/3))
else:
#construct map from residue to anm property and dx,dy,dz vectors
propmap = []
for (i,a) in enumerate(atoms.calpha):
propmap.append({'chain': a.getChid(), 'resi':a.getResnum(),
'props': {'dy': aarr[3*i+1], 'dz': aarr[3*i+2] } });
#set the atom property
#TODO: implement something more efficient on the 3Dmol.js side (this is O(n*m)!)
view.mapAtomProperties(propmap)
#create vibrations
frames = kwargs.get('frames',10)
amplitude = kwargs.get('amplitude',100)
view.vibrate(frames, amplitude)
animate = kwargs.get('animate',{'loop':'rock'})
view.animate(animate)
if 'style' in kwargs: # this is never a list
view.setStyle({},kwargs['style'])
if 'styles' in kwargs:
#allow simpler forms - convert them into a list
styles = kwargs['styles']
if type(styles) == dict:
styles = ({},styles)
if type(styles) == tuple:
styles = [styles]
for (sel, style) in styles:
view.setStyle(sel, style)
zoomto = kwargs['zoomto'] if 'zoomto' in kwargs else kwargs.get('zoomTo',{})
view.zoomTo(zoomto)
return view
