def display_molecule(molecule, highlight=None, **kwargs):
topology = {
'atom_types': molecule.type_array,
'bonds': molecule.bonds,
'atom_names': molecule.atom_name,
'secondary_structure': molecule.secondary_structure
}
mv = MolecularViewer(molecule.r_array.astype('float32'), topology)
kind = kwargs.get('kind', 'wireframe')
if kind == 'wireframe':
if molecule.n_bonds != 0:
mv.points(size=0.15, highlight=highlight)
mv.lines()
else:
mv.points(highlight=highlight)
elif kind == 'ball_and_sticks':
mv.ball_and_sticks()
elif kind == 'cartoon':
mv.cartoon()
else:
raise ValueError("kind {} not found".format(kind))
return mv
def display_molecule(molecule):
topology = {'atom_types': molecule.type_array, 'bonds': molecule.bonds}
mv = MolecularViewer(molecule.r_array, topology)
if molecule.n_bonds != 0:
mv.points(size=0.15)
mv.lines()
else:
mv.points()
return mv
def view(a,
colour=None,
bonds=True,
cell=True,
scale=10.0,
cutoff_scale=1.2,
cmap=None,
vmin=None,
vmax=None):
topology = {}
topology['atom_types'] = a.get_chemical_symbols()
if bonds:
n = a.numbers
maxn = n.max()
cutoffs = np.zeros([maxn + 1, maxn + 1])
for n1, n2 in itertools.product(n, n):
cutoffs[n1, n2] = cutoff_scale * (covalent_radii[n1] +
covalent_radii[n2])
# Construct a bond list
i, j, S = neighbour_list('ijS', a, cutoffs,
np.array(a.numbers, dtype=np.int32))
m = np.logical_and(i < j, (S == 0).all(axis=1))
i = i[m]
j = j[m]
topology['bonds'] = [(x, y) for x, y in zip(i, j)]
colorlist = None
if colour is not None:
colour = np.array(colour, dtype=np.float64)
if cmap is None:
from matplotlib.cm import jet
cmap = jet
if vmin is None:
vmin = np.min(colour)
if vmax is None:
vmax = np.max(colour)
colour = (colour - vmin) / (vmax - vmin)
colorlist = [
'0x%02x%02x%02x' % (r * 256, g * 256, b * 256)
for (r, g, b, alpha) in cmap(colour)
]
mv = MolecularViewer(a.positions / scale, topology=topology)
mv.ball_and_sticks(colorlist=colorlist)
if cell:
O = np.zeros(3, dtype=np.float32)
La, Lb, Lc = a.cell.astype(np.float32) / scale
start = np.r_[O, O, O, O + Lb, O + Lc, O + La, O + Lc, O + La, O + Lb,
O + Lb + Lc, O + La + Lc, O + La + Lb]
end = np.r_[O + La, O + Lb, O + Lc, O + Lb + La, O + Lc + Lb,
O + La + Lc, O + Lc + La, O + La + Lb, O + Lb + Lc,
O + Lb + Lc + La, O + La + Lc + Lb, O + La + Lb + Lc]
rgb = [0xFF0000, 0x00FF00, 0x0000FF] * 4
mv.add_representation(
'lines', {
'startCoords': start,
'endCoords': end,
'startColors': rgb,
'endColors': rgb
})
return mv
def quick_view(structure,
bonds=True,
conventional=False,
transform=None,
show_box=True,
bond_tol=0.2,
stick_radius=0.1):
"""
A function to visualize pymatgen Structure objects in jupyter notebook using chemview package.
Args:
structure: pymatgen Structure
bonds: (bool) visualize bonds. Bonds are found by comparing distances
to added covalent radii of pairs. Defaults to True.
conventional: (bool) use conventional cell. Defaults to False.
transform: (list) can be used to make supercells with pymatgen.Structure.make_supercell method
show_box: (bool) unit cell is shown. Defaults to True.
bond_tol: (float) used if bonds=True. Sets the extra distance tolerance when finding bonds.
stick_radius: (float) radius of bonds.
Returns:
A chemview.MolecularViewer object
"""
s = structure.copy()
if conventional:
s = SpacegroupAnalyzer(s).get_conventional_standard_structure()
if transform:
s.make_supercell(transform)
atom_types = [i.symbol for i in s.species]
if bonds:
bonds = []
for i in range(s.num_sites - 1):
sym_i = s[i].specie.symbol
for j in range(i + 1, s.num_sites):
sym_j = s[j].specie.symbol
max_d = CovalentRadius.radius[sym_i] + CovalentRadius.radius[
sym_j] + bond_tol
if s.get_distance(i, j, np.array([0, 0, 0])) < max_d:
bonds.append((i, j))
bonds = bonds if bonds else None
mv = MolecularViewer(s.cart_coords,
topology={
'atom_types': atom_types,
'bonds': bonds
})
if bonds:
mv.ball_and_sticks(stick_radius=stick_radius)
for i in s.sites:
el = i.specie.symbol
coord = i.coords
r = CovalentRadius.radius[el]
mv.add_representation(
'spheres', {
'coordinates': coord.astype('float32'),
'colors': [get_atom_color(el)],
'radii': [r * 0.5],
'opacity': 1.0
})
if show_box:
o = np.array([0, 0, 0])
a, b, c = s.lattice.matrix[0], s.lattice.matrix[1], s.lattice.matrix[2]
starts = [o, o, o, a, a, b, b, c, c, a + b, a + c, b + c]
ends = [
a, b, c, a + b, a + c, b + a, b + c, c + a, c + b, a + b + c,
a + b + c, a + b + c
]
colors = [0xffffff for i in range(12)]
mv.add_representation(
'lines', {
'startCoords': np.array(starts),
'endCoords': np.array(ends),
'startColors': colors,
'endColors': colors
})
return mv