index | modules | gitlab

# Source code for ase.visualize.ngl

    from ase import Atoms

    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, IntSlider, HBox, 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()

        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()

    def view_ngl(atoms, 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
        """
        return NGLDisplay(atoms, w, h).gui