Coverage for /builds/kinetik161/ase/ase/atom.py: 94.44%

1"""This module defines the Atom object.""" 

2 

3import numpy as np 

4 

5from ase.data import atomic_masses, atomic_numbers, chemical_symbols 

6 

7# Singular, plural, default value: 

8names = {'position': ('positions', np.zeros(3)), 

9 'number': ('numbers', 0), 

10 'tag': ('tags', 0), 

11 'momentum': ('momenta', np.zeros(3)), 

12 'mass': ('masses', None), 

13 'magmom': ('initial_magmoms', 0.0), 

14 'charge': ('initial_charges', 0.0)} 

15 

16 

17def atomproperty(name, doc): 

18 """Helper function to easily create Atom attribute property.""" 

19 

20 def getter(self): 

21 return self.get(name) 

22 

23 def setter(self, value): 

24 self.set(name, value) 

25 

26 def deleter(self): 

27 self.delete(name) 

28 

29 return property(getter, setter, deleter, doc) 

30 

31 

32def abcproperty(index): 

33 """Helper function to easily create Atom ABC-property.""" 

34 

35 def getter(self): 

36 return self.scaled_position[index] 

37 

38 def setter(self, value): 

39 # We can't just do self.scaled_position[i] = value 

40 # because scaled_position is a new buffer, not a view into 

41 # something we can write back to. 

42 # This is a clear bug! 

43 spos = self.scaled_position 

44 spos[index] = value 

45 self.scaled_position = spos 

46 

47 return property(getter, setter, doc='ABC'[index] + '-coordinate') 

48 

49 

50def xyzproperty(index): 

51 """Helper function to easily create Atom XYZ-property.""" 

52 

53 def getter(self): 

54 return self.position[index] 

55 

56 def setter(self, value): 

57 self.position[index] = value 

58 

59 return property(getter, setter, doc='XYZ'[index] + '-coordinate') 

60 

61 

62class Atom: 

63 """Class for representing a single atom. 

64 

65 Parameters: 

66 

67 symbol: str or int 

68 Can be a chemical symbol (str) or an atomic number (int). 

69 position: sequence of 3 floats 

70 Atomic position. 

71 tag: int 

72 Special purpose tag. 

73 momentum: sequence of 3 floats 

74 Momentum for atom. 

75 mass: float 

76 Atomic mass in atomic units. 

77 magmom: float or 3 floats 

78 Magnetic moment. 

79 charge: float 

80 Atomic charge. 

81 """ 

82 __slots__ = ['data', 'atoms', 'index'] 

83 

84 def __init__(self, symbol='X', position=(0, 0, 0), 

85 tag=None, momentum=None, mass=None, 

86 magmom=None, charge=None, 

87 atoms=None, index=None): 

88 

89 self.data = d = {} 

90 

91 if atoms is None: 

92 # This atom is not part of any Atoms object: 

93 if isinstance(symbol, str): 

94 d['number'] = atomic_numbers[symbol] 

95 else: 

96 d['number'] = symbol 

97 d['position'] = np.array(position, float) 

98 d['tag'] = tag 

99 if momentum is not None: 

100 momentum = np.array(momentum, float) 

101 d['momentum'] = momentum 

102 d['mass'] = mass 

103 if magmom is not None: 

104 magmom = np.array(magmom, float) 

105 d['magmom'] = magmom 

106 d['charge'] = charge 

107 

108 self.index = index 

109 self.atoms = atoms 

110 

111 @property 

112 def scaled_position(self): 

113 pos = self.position 

114 spos = self.atoms.cell.scaled_positions(pos[np.newaxis]) 

115 return spos[0] 

116 

117 @scaled_position.setter 

118 def scaled_position(self, value): 

119 pos = self.atoms.cell.cartesian_positions(value) 

120 self.position = pos 

121 

122 def __repr__(self): 

123 s = f"Atom('{self.symbol}', {list(self.position)}" 

124 for name in ['tag', 'momentum', 'mass', 'magmom', 'charge']: 

125 value = self.get_raw(name) 

126 if value is not None: 

127 if isinstance(value, np.ndarray): 

128 value = value.tolist() 

129 s += f', {name}={value}' 

130 if self.atoms is None: 

131 s += ')' 

132 else: 

133 s += ', index=%d)' % self.index 

134 return s 

135 

136 def cut_reference_to_atoms(self): 

137 """Cut reference to atoms object.""" 

138 for name in names: 

139 self.data[name] = self.get_raw(name) 

140 self.index = None 

141 self.atoms = None 

142 

143 def get_raw(self, name): 

144 """Get name attribute, return None if not explicitly set.""" 

145 if name == 'symbol': 

146 return chemical_symbols[self.get_raw('number')] 

147 

148 if self.atoms is None: 

149 return self.data[name] 

150 

151 plural = names[name][0] 

152 if plural in self.atoms.arrays: 

153 return self.atoms.arrays[plural][self.index] 

154 else: 

155 return None 

156 

157 def get(self, name): 

158 """Get name attribute, return default if not explicitly set.""" 

159 value = self.get_raw(name) 

160 if value is None: 

161 if name == 'mass': 

162 value = atomic_masses[self.number] 

163 else: 

164 value = names[name][1] 

165 return value 

166 

167 def set(self, name, value): 

168 """Set name attribute to value.""" 

169 if name == 'symbol': 

170 name = 'number' 

171 value = atomic_numbers[value] 

172 

173 if self.atoms is None: 

174 assert name in names 

175 self.data[name] = value 

176 else: 

177 plural, default = names[name] 

178 if plural in self.atoms.arrays: 

179 array = self.atoms.arrays[plural] 

180 if name == 'magmom' and array.ndim == 2: 

181 assert len(value) == 3 

182 array[self.index] = value 

183 else: 

184 if name == 'magmom' and np.asarray(value).ndim == 1: 

185 array = np.zeros((len(self.atoms), 3)) 

186 elif name == 'mass': 

187 array = self.atoms.get_masses() 

188 else: 

189 default = np.asarray(default) 

190 array = np.zeros((len(self.atoms),) + default.shape, 

191 default.dtype) 

192 array[self.index] = value 

193 self.atoms.new_array(plural, array) 

194 

195 def delete(self, name): 

196 """Delete name attribute.""" 

197 assert self.atoms is None 

198 assert name not in ['number', 'symbol', 'position'] 

199 self.data[name] = None 

200 

201 symbol = atomproperty('symbol', 'Chemical symbol') 

202 number = atomproperty('number', 'Atomic number') 

203 position = atomproperty('position', 'XYZ-coordinates') 

204 tag = atomproperty('tag', 'Integer tag') 

205 momentum = atomproperty('momentum', 'XYZ-momentum') 

206 mass = atomproperty('mass', 'Atomic mass') 

207 magmom = atomproperty('magmom', 'Initial magnetic moment') 

208 charge = atomproperty('charge', 'Initial atomic charge') 

209 x = xyzproperty(0) 

210 y = xyzproperty(1) 

211 z = xyzproperty(2) 

212 

213 a = abcproperty(0) 

214 b = abcproperty(1) 

215 c = abcproperty(2)