1"""colors.py - select how to color the atoms in the GUI.""" 

2from ase.gui.i18n import _ 

3 

4import numpy as np 

5 

6import ase.gui.ui as ui 

7from ase.gui.utils import get_magmoms 

8 

9 

10class ColorWindow: 

11 """A window for selecting how to color the atoms.""" 

12 

13 def __init__(self, gui): 

14 self.reset(gui) 

15 

16 def reset(self, gui): 

17 """create a new color window""" 

18 self.win = ui.Window(_('Colors'), wmtype='utility') 

19 self.gui = gui 

20 self.win.add(ui.Label(_('Choose how the atoms are colored:'))) 

21 values = ['jmol', 'tag', 'force', 'velocity', 

22 'initial charge', 'magmom', 'neighbors'] 

23 labels = [_('By atomic number, default "jmol" colors'), 

24 _('By tag'), 

25 _('By force'), 

26 _('By velocity'), 

27 _('By initial charge'), 

28 _('By magnetic moment'), 

29 _('By number of neighbors'), ] 

30 

31 haveit = ['numbers', 'positions', 'forces', 'momenta', 

32 'initial_charges', 'initial_magmoms'] 

33 for key in self.gui.atoms.arrays: 

34 if key not in haveit: 

35 values.append(key) 

36 labels.append('By user-defined "{}"'.format(key)) 

37 

38 self.radio = ui.RadioButtons(labels, values, self.toggle, 

39 vertical=True) 

40 self.radio.value = gui.colormode 

41 self.win.add(self.radio) 

42 self.activate() 

43 self.label = ui.Label() 

44 self.win.add(self.label) 

45 

46 if hasattr(self, 'mnmx'): 

47 self.win.add(self.cmaps) 

48 self.win.add(self.mnmx) 

49 

50 def change_mnmx(self, mn=None, mx=None): 

51 """change min and/or max values for colormap""" 

52 if mn: 

53 self.mnmx[1].value = mn 

54 if mx: 

55 self.mnmx[3].value = mx 

56 mn, mx = self.mnmx[1].value, self.mnmx[3].value 

57 colorscale, _, _ = self.gui.colormode_data 

58 self.gui.colormode_data = colorscale, mn, mx 

59 self.gui.draw() 

60 

61 def activate(self): 

62 images = self.gui.images 

63 atoms = self.gui.atoms 

64 radio = self.radio 

65 radio['tag'].active = atoms.has('tags') 

66 

67 # XXX not sure how to deal with some images having forces, 

68 # and other images not. Same goes for below quantities 

69 F = images.get_forces(atoms) 

70 radio['force'].active = F is not None 

71 radio['velocity'].active = atoms.has('momenta') 

72 radio['initial charge'].active = atoms.has('initial_charges') 

73 radio['magmom'].active = get_magmoms(atoms).any() 

74 radio['neighbors'].active = True 

75 

76 def toggle(self, value): 

77 self.gui.colormode = value 

78 if value == 'jmol' or value == 'neighbors': 

79 if hasattr(self, 'mnmx'): 

80 "delete the min max fields by creating a new window" 

81 del self.mnmx 

82 del self.cmaps 

83 self.win.close() 

84 self.reset(self.gui) 

85 text = '' 

86 else: 

87 scalars = np.ma.array([self.gui.get_color_scalars(i) 

88 for i in range(len(self.gui.images))]) 

89 mn = np.min(scalars) 

90 mx = np.max(scalars) 

91 self.gui.colormode_data = None, mn, mx 

92 

93 cmaps = ['default', 'old'] 

94 try: 

95 import pylab as plt 

96 cmaps += [m for m in plt.cm.datad if not m.endswith("_r")] 

97 except ImportError: 

98 pass 

99 self.cmaps = [_('cmap:'), 

100 ui.ComboBox(cmaps, cmaps, self.update_colormap), 

101 _('N:'), 

102 ui.SpinBox(26, 0, 100, 1, self.update_colormap)] 

103 self.update_colormap('default') 

104 

105 try: 

106 unit = {'tag': '', 

107 'force': 'eV/Ang', 

108 'velocity': '(eV/amu)^(1/2)', 

109 'charge': '|e|', 

110 'initial charge': '|e|', 

111 u'magmom': 'μB'}[value] 

112 except KeyError: 

113 unit = '' 

114 text = '' 

115 

116 rng = mx - mn # XXX what are optimal allowed range and steps ? 

117 self.mnmx = [_('min:'), 

118 ui.SpinBox(mn, mn - 10 * rng, mx + rng, rng / 10., 

119 self.change_mnmx, width=20), 

120 _('max:'), 

121 ui.SpinBox(mx, mn - 10 * rng, mx + rng, rng / 10., 

122 self.change_mnmx, width=20), 

123 _(unit)] 

124 self.win.close() 

125 self.reset(self.gui) 

126 

127 self.label.text = text 

128 self.radio.value = value 

129 self.gui.draw() 

130 return text # for testing 

131 

132 def notify_atoms_changed(self): 

133 "Called by gui object when the atoms have changed." 

134 self.activate() 

135 mode = self.gui.colormode 

136 if not self.radio[mode].active: 

137 mode = 'jmol' 

138 self.toggle(mode) 

139 

140 def update_colormap(self, cmap=None, N=26): 

141 "Called by gui when colormap has changed" 

142 if cmap is None: 

143 cmap = self.cmaps[1].value 

144 try: 

145 N = int(self.cmaps[3].value) 

146 except AttributeError: 

147 N = 26 

148 colorscale, mn, mx = self.gui.colormode_data 

149 if cmap == 'default': 

150 colorscale = ['#{0:02X}80{0:02X}'.format(int(red)) 

151 for red in np.linspace(0, 250, N)] 

152 elif cmap == 'old': 

153 colorscale = ['#{0:02X}AA00'.format(int(red)) 

154 for red in np.linspace(0, 230, N)] 

155 else: 

156 try: 

157 import pylab as plt 

158 import matplotlib 

159 cmap = plt.cm.get_cmap(cmap) 

160 colorscale = [matplotlib.colors.rgb2hex(c[:3]) for c in 

161 cmap(np.linspace(0, 1, N))] 

162 except (ImportError, ValueError) as e: 

163 raise RuntimeError('Can not load colormap {0}: {1}'.format( 

164 cmap, str(e))) 

165 self.gui.colormode_data = colorscale, mn, mx 

166 self.gui.draw()