from ase import Atoms
[docs]class NGLDisplay:
"""Structure display class
Provides basic structure/trajectory display
in the notebook and optional gui which can be used to enhance its
usability. It is also possible to extend the functionality of the
particular instance of the viewer by adding further widgets
manipulating the structure.
"""
def __init__(self, atoms, xsize=500, ysize=500):
import nglview
import nglview.color
from ipywidgets import Dropdown, FloatSlider, HBox, IntSlider, VBox
self.atoms = atoms
if isinstance(atoms[0], Atoms):
# Assume this is a trajectory or struct list
self.view = nglview.show_asetraj(atoms, default=False)
self.frm = IntSlider(value=0, min=0, max=len(atoms) - 1)
self.frm.observe(self._update_frame)
self.struct = atoms[0]
else:
# Assume this is just a single structure
self.view = nglview.show_ase(atoms, default=False)
self.struct = atoms
self.frm = None
self.colors = {}
self.view._remote_call('setSize', target='Widget',
args=['%dpx' % (xsize,), '%dpx' % (ysize,)])
self.view.add_unitcell()
self.view.add_spacefill()
self.view.camera = 'orthographic'
self.view.parameters = {"clipDist": 0}
self.view.center()
self.asel = Dropdown(options=['All'] +
list(set(self.struct.get_chemical_symbols())),
value='All', description='Show')
self.csel = Dropdown(options=nglview.color.COLOR_SCHEMES,
value='element', description='Color scheme')
self.rad = FloatSlider(value=0.5, min=0.0, max=1.5, step=0.01,
description='Ball size')
self.asel.observe(self._select_atom)
self.csel.observe(self._update_repr)
self.rad.observe(self._update_repr)
self.view.update_spacefill(radiusType='covalent',
radiusScale=0.5,
color_scheme=self.csel.value,
color_scale='rainbow')
wdg = [self.asel, self.csel, self.rad]
if self.frm:
wdg.append(self.frm)
self.gui = HBox([self.view, VBox(wdg)])
# Make useful shortcuts for the user of the class
self.gui.view = self.view
self.gui.control_box = self.gui.children[1]
self.gui.custom_colors = self.custom_colors
def _update_repr(self, chg=None):
self.view.update_spacefill(radiusType='covalent',
radiusScale=self.rad.value,
color_scheme=self.csel.value,
color_scale='rainbow')
def _update_frame(self, chg=None):
self.view.frame = self.frm.value
return
def _select_atom(self, chg=None):
sel = self.asel.value
self.view.remove_spacefill()
for e in set(self.struct.get_chemical_symbols()):
if (sel == 'All' or e == sel):
if e in self.colors:
self.view.add_spacefill(selection='#' + e,
color=self.colors[e])
else:
self.view.add_spacefill(selection='#' + e)
self._update_repr()
[docs] def custom_colors(self, clr=None):
"""
Define custom colors for some atoms. Pass a dictionary of the form
{'Fe':'red', 'Au':'yellow'} to the function.
To reset the map to default call the method without parameters.
"""
if clr:
self.colors = clr
else:
self.colors = {}
self._select_atom()
[docs]def view_ngl(atoms, data=None, repeat=None, w=500, h=500):
"""
Returns the nglviewer + some control widgets in the VBox ipywidget.
The viewer supports any Atoms objectand any sequence of Atoms objects.
The returned object has two shortcuts members:
.view:
nglviewer ipywidget for direct interaction
.control_box:
VBox ipywidget containing view control widgets
"""
# TODO: make `data` and `repeat` really available
return NGLDisplay(atoms, w, h).gui