1from ase import Atoms 

2 

3 

4class NGLDisplay: 

5 """Structure display class 

6 

7 Provides basic structure/trajectory display 

8 in the notebook and optional gui which can be used to enhance its 

9 usability. It is also possible to extend the functionality of the 

10 particular instance of the viewer by adding further widgets 

11 manipulating the structure. 

12 """ 

13 def __init__(self, atoms, xsize=500, ysize=500): 

14 import nglview 

15 import nglview.color 

16 

17 from ipywidgets import Dropdown, FloatSlider, IntSlider, HBox, VBox 

18 self.atoms = atoms 

19 if isinstance(atoms[0], Atoms): 

20 # Assume this is a trajectory or struct list 

21 self.view = nglview.show_asetraj(atoms, default=False) 

22 self.frm = IntSlider(value=0, min=0, max=len(atoms) - 1) 

23 self.frm.observe(self._update_frame) 

24 self.struct = atoms[0] 

25 else: 

26 # Assume this is just a single structure 

27 self.view = nglview.show_ase(atoms, default=False) 

28 self.struct = atoms 

29 self.frm = None 

30 

31 self.colors = {} 

32 self.view._remote_call('setSize', target='Widget', 

33 args=['%dpx' % (xsize,), '%dpx' % (ysize,)]) 

34 self.view.add_unitcell() 

35 self.view.add_spacefill() 

36 self.view.camera = 'orthographic' 

37 self.view.parameters = {"clipDist": 0} 

38 

39 self.view.center() 

40 

41 self.asel = Dropdown(options=['All'] + 

42 list(set(self.struct.get_chemical_symbols())), 

43 value='All', description='Show') 

44 

45 self.csel = Dropdown(options=nglview.color.COLOR_SCHEMES, 

46 value='element', description='Color scheme') 

47 

48 self.rad = FloatSlider(value=0.5, min=0.0, max=1.5, step=0.01, 

49 description='Ball size') 

50 

51 self.asel.observe(self._select_atom) 

52 self.csel.observe(self._update_repr) 

53 self.rad.observe(self._update_repr) 

54 

55 self.view.update_spacefill(radiusType='covalent', 

56 radiusScale=0.5, 

57 color_scheme=self.csel.value, 

58 color_scale='rainbow') 

59 

60 wdg = [self.asel, self.csel, self.rad] 

61 if self.frm: 

62 wdg.append(self.frm) 

63 

64 self.gui = HBox([self.view, VBox(wdg)]) 

65 # Make useful shortcuts for the user of the class 

66 self.gui.view = self.view 

67 self.gui.control_box = self.gui.children[1] 

68 self.gui.custom_colors = self.custom_colors 

69 

70 def _update_repr(self, chg=None): 

71 self.view.update_spacefill(radiusType='covalent', 

72 radiusScale=self.rad.value, 

73 color_scheme=self.csel.value, 

74 color_scale='rainbow') 

75 

76 def _update_frame(self, chg=None): 

77 self.view.frame = self.frm.value 

78 return 

79 

80 def _select_atom(self, chg=None): 

81 sel = self.asel.value 

82 self.view.remove_spacefill() 

83 for e in set(self.struct.get_chemical_symbols()): 

84 if (sel == 'All' or e == sel): 

85 if e in self.colors: 

86 self.view.add_spacefill(selection='#' + e, 

87 color=self.colors[e]) 

88 else: 

89 self.view.add_spacefill(selection='#' + e) 

90 self._update_repr() 

91 

92 def custom_colors(self, clr=None): 

93 """ 

94 Define custom colors for some atoms. Pass a dictionary of the form 

95 {'Fe':'red', 'Au':'yellow'} to the function. 

96 To reset the map to default call the method without parameters. 

97 """ 

98 if clr: 

99 self.colors = clr 

100 else: 

101 self.colors = {} 

102 self._select_atom() 

103 

104 

105def view_ngl(atoms, w=500, h=500): 

106 """ 

107 Returns the nglviewer + some control widgets in the VBox ipywidget. 

108 The viewer supports any Atoms objectand any sequence of Atoms objects. 

109 The returned object has two shortcuts members: 

110 

111 .view: 

112 nglviewer ipywidget for direct interaction 

113 .control_box: 

114 VBox ipywidget containing view control widgets 

115 """ 

116 return NGLDisplay(atoms, w, h).gui