Coverage for /builds/ase/ase/ase/gui/colors.py: 81.90%

1# fmt: off 

2 

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

4import numpy as np 

5 

6import ase.gui.ui as ui 

7from ase.gui.i18n import _ 

8from ase.gui.utils import get_magmoms 

9 

10 

11class ColorWindow: 

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

13 

14 def __init__(self, gui): 

15 self.reset(gui) 

16 

17 def reset(self, gui): 

18 """create a new color window""" 

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

20 self.gui = gui 

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

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

23 'initial charge', 'magmom', 'neighbors'] 

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

25 _('By tag'), 

26 _('By force'), 

27 _('By velocity'), 

28 _('By initial charge'), 

29 _('By magnetic moment'), 

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

31 

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

33 'initial_charges', 'initial_magmoms'] 

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

35 if key not in haveit: 

36 values.append(key) 

37 labels.append(f'By user-defined "{key}"') 

38 

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

40 vertical=True) 

41 self.radio.value = gui.colormode 

42 self.win.add(self.radio) 

43 self.activate() 

44 self.label = ui.Label() 

45 self.win.add(self.label) 

46 

47 if hasattr(self, 'mnmx'): 

48 self.win.add(self.cmaps) 

49 self.win.add(self.mnmx) 

50 

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

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

53 if mn: 

54 self.mnmx[1].value = mn 

55 if mx: 

56 self.mnmx[3].value = mx 

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

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

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

60 self.gui.draw() 

61 

62 def activate(self): 

63 images = self.gui.images 

64 atoms = self.gui.atoms 

65 radio = self.radio 

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

67 

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

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

70 F = images.get_forces(atoms) 

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

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

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

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

75 radio['neighbors'].active = True 

76 

77 def toggle(self, value): 

78 self.gui.colormode = value 

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

80 if hasattr(self, 'mnmx'): 

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

82 del self.mnmx 

83 del self.cmaps 

84 self.win.close() 

85 self.reset(self.gui) 

86 text = '' 

87 else: 

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

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

90 mn = np.min(scalars) 

91 mx = np.max(scalars) 

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

93 

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

95 try: 

96 import pylab as plt 

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

98 except ImportError: 

99 pass 

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

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

102 _('N:'), 

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

104 self.update_colormap('default') 

105 

106 try: 

107 unit = {'tag': '', 

108 'force': 'eV/Ang', 

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

110 'charge': '|e|', 

111 'initial charge': '|e|', 

112 'magmom': 'μB'}[value] 

113 except KeyError: 

114 unit = '' 

115 text = '' 

116 

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

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

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

120 self.change_mnmx, width=20), 

121 _('max:'), 

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

123 self.change_mnmx, width=20), 

124 _(unit)] 

125 self.win.close() 

126 self.reset(self.gui) 

127 

128 self.label.text = text 

129 self.radio.value = value 

130 self.gui.draw() 

131 return text # for testing 

132 

133 def notify_atoms_changed(self): 

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

135 self.activate() 

136 mode = self.gui.colormode 

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

138 mode = 'jmol' 

139 self.toggle(mode) 

140 

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

142 "Called by gui when colormap has changed" 

143 import matplotlib 

144 if cmap is None: 

145 cmap = self.cmaps[1].value 

146 try: 

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

148 except AttributeError: 

149 N = 26 

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

151 if cmap == 'default': 

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

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

154 elif cmap == 'old': 

155 colorscale = [f'#{int(red):02X}AA00' 

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

157 else: 

158 cmap_obj = matplotlib.colormaps[cmap] 

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

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

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

162 self.gui.draw()