Coverage for /builds/ase/ase/ase/visualize/ngl.py: 14.04%

1# fmt: off 

2 

3from ase import Atoms 

4 

5 

6class NGLDisplay: 

7 """Structure display class 

8 

9 Provides basic structure/trajectory display 

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

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

12 particular instance of the viewer by adding further widgets 

13 manipulating the structure. 

14 """ 

15 

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

17 import nglview 

18 import nglview.color 

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

20 self.atoms = atoms 

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

22 # Assume this is a trajectory or struct list 

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

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

25 self.frm.observe(self._update_frame) 

26 self.struct = atoms[0] 

27 else: 

28 # Assume this is just a single structure 

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

30 self.struct = atoms 

31 self.frm = None 

32 

33 self.colors = {} 

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

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

36 self.view.add_unitcell() 

37 self.view.add_spacefill() 

38 self.view.camera = 'orthographic' 

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

40 

41 self.view.center() 

42 

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

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

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

46 

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

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

49 

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

51 description='Ball size') 

52 

53 self.asel.observe(self._select_atom) 

54 self.csel.observe(self._update_repr) 

55 self.rad.observe(self._update_repr) 

56 

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

58 radiusScale=0.5, 

59 color_scheme=self.csel.value, 

60 color_scale='rainbow') 

61 

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

63 if self.frm: 

64 wdg.append(self.frm) 

65 

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

67 # Make useful shortcuts for the user of the class 

68 self.gui.view = self.view 

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

70 self.gui.custom_colors = self.custom_colors 

71 

72 def _update_repr(self, chg=None): 

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

74 radiusScale=self.rad.value, 

75 color_scheme=self.csel.value, 

76 color_scale='rainbow') 

77 

78 def _update_frame(self, chg=None): 

79 self.view.frame = self.frm.value 

80 return 

81 

82 def _select_atom(self, chg=None): 

83 sel = self.asel.value 

84 self.view.remove_spacefill() 

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

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

87 if e in self.colors: 

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

89 color=self.colors[e]) 

90 else: 

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

92 self._update_repr() 

93 

94 def custom_colors(self, clr=None): 

95 """ 

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

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

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

99 """ 

100 if clr: 

101 self.colors = clr 

102 else: 

103 self.colors = {} 

104 self._select_atom() 

105 

106 

107def view_ngl(atoms, data=None, repeat=None, w=500, h=500): 

108 """ 

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

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

111 The returned object has two shortcuts members: 

112 

113 .view: 

114 nglviewer ipywidget for direct interaction 

115 .control_box: 

116 VBox ipywidget containing view control widgets 

117 """ 

118 # TODO: make `data` and `repeat` really available 

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