geometry_viewer.py 文件源码

def set_color(self, colors, atoms=None, save=True):
        """ Set atom colors

        May be called in several different ways:
          - ``set_color(color, atoms=list_of_atoms_or_None)``
                  where all passed atoms are to be colored a single color
          - ``set_color(list_of_colors, atoms=list_of_atoms_or_None)``
                  with a list of colors for each atom
          -  ``set_color(dict_from_atoms_to_colors)``
                  a dictionary that maps atoms to colors
          - ``set_color(f, atoms=list_of_atoms_or_None)``
                 where f is a function that maps atoms to colors

        Args:
            colors (see note for allowable types): list of colors for each atom, or map
               from atoms to colors, or a single color for all atoms
            atoms (List[moldesign.Atom]): list of atoms (if None, assumed to be mol.atoms; ignored
               if a dict is passed for "color")
            save (bool): always color these atoms this way (until self.unset_color is called)

        See Also:
            :method:`GeometryViewer.color_by`` - to automatically color atoms using numerical
               and categorical data
        """
        if hasattr(colors, 'items'):
            atoms, colors = zip(*colors.items())
        elif atoms is None:
            atoms = self.mol.atoms

        if callable(colors):
            colors = map(colors, atoms)
        elif isinstance(colors, basestring) or not hasattr(colors, '__iter__'):
            colors = [colors for atom in atoms]

        for atom,color in zip(atoms, colors):
            c = translate_color(color, '#')
            if save:
                self.atom_colors[atom] = c
            self.styles[str(atom.index)]['color'] = c
        self.send_state('styles')