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 def __init__(self, gui): 

13 self.reset(gui) 

14 

15 def reset(self, gui): 

16 """create a new color window""" 

17 self.win = ui.Window(_('Colors')) 

18 self.gui = gui 

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

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

21 'initial charge', 'magmom', 'neighbors'] 

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

23 _('By tag'), 

24 _('By force'), 

25 _('By velocity'), 

26 _('By initial charge'), 

27 _('By magnetic moment'), 

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

29 

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

31 'initial_charges', 'initial_magmoms'] 

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

33 if key not in haveit: 

34 values.append(key) 

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

36 

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

38 vertical=True) 

39 self.radio.value = gui.colormode 

40 self.win.add(self.radio) 

41 self.activate() 

42 self.label = ui.Label() 

43 self.win.add(self.label) 

44 

45 if hasattr(self, 'mnmx'): 

46 self.win.add(self.cmaps) 

47 self.win.add(self.mnmx) 

48 

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

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

51 if mn: 

52 self.mnmx[1].value = mn 

53 if mx: 

54 self.mnmx[3].value = mx 

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

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

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

58 self.gui.draw() 

59 

60 def activate(self): 

61 images = self.gui.images 

62 atoms = self.gui.atoms 

63 radio = self.radio 

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

65 

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

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

68 F = images.get_forces(atoms) 

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

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

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

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

73 radio['neighbors'].active = True 

74 

75 def toggle(self, value): 

76 self.gui.colormode = value 

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

78 if hasattr(self, 'mnmx'): 

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

80 del self.mnmx 

81 del self.cmaps 

82 self.win.close() 

83 self.reset(self.gui) 

84 text = '' 

85 else: 

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

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

88 mn = np.min(scalars) 

89 mx = np.max(scalars) 

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

91 

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

93 try: 

94 import pylab as plt 

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

96 except ImportError: 

97 pass 

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

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

100 _('N:'), 

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

102 self.update_colormap('default') 

103 

104 try: 

105 unit = {'tag': '', 

106 'force': 'eV/Ang', 

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

108 'charge': '|e|', 

109 'initial charge': '|e|', 

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

111 except KeyError: 

112 unit = '' 

113 text = '' 

114 

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

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

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

118 self.change_mnmx, width=20), 

119 _('max:'), 

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

121 self.change_mnmx, width=20), 

122 _(unit)] 

123 self.win.close() 

124 self.reset(self.gui) 

125 

126 self.label.text = text 

127 self.radio.value = value 

128 self.gui.draw() 

129 return text # for testing 

130 

131 def notify_atoms_changed(self): 

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

133 self.activate() 

134 mode = self.gui.colormode 

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

136 mode = 'jmol' 

137 self.toggle(mode) 

138 

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

140 "Called by gui when colormap has changed" 

141 if cmap is None: 

142 cmap = self.cmaps[1].value 

143 try: 

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

145 except AttributeError: 

146 N = 26 

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

148 if cmap == 'default': 

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

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

151 elif cmap == 'old': 

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

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

154 else: 

155 try: 

156 import pylab as plt 

157 import matplotlib 

158 cmap = plt.cm.get_cmap(cmap) 

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

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

161 except (ImportError, ValueError) as e: 

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

163 cmap, str(e))) 

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

165 self.gui.draw()