1from random import randint 

2from typing import Dict, Any 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.constraints import dict2constraint 

8from ase.calculators.calculator import (all_properties, 

9 PropertyNotImplementedError, 

10 kptdensity2monkhorstpack) 

11from ase.calculators.singlepoint import SinglePointCalculator 

12from ase.data import chemical_symbols, atomic_masses 

13from ase.formula import Formula 

14from ase.geometry import cell_to_cellpar 

15from ase.io.jsonio import decode 

16 

17 

18class FancyDict(dict): 

19 """Dictionary with keys available as attributes also.""" 

20 def __getattr__(self, key): 

21 if key not in self: 

22 return dict.__getattribute__(self, key) 

23 value = self[key] 

24 if isinstance(value, dict): 

25 return FancyDict(value) 

26 return value 

27 

28 def __dir__(self): 

29 return self.keys() # for tab-completion 

30 

31 

32def atoms2dict(atoms): 

33 dct = { 

34 'numbers': atoms.numbers, 

35 'positions': atoms.positions, 

36 'unique_id': '%x' % randint(16**31, 16**32 - 1)} 

37 if atoms.pbc.any(): 

38 dct['pbc'] = atoms.pbc 

39 if atoms.cell.any(): 

40 dct['cell'] = atoms.cell 

41 if atoms.has('initial_magmoms'): 

42 dct['initial_magmoms'] = atoms.get_initial_magnetic_moments() 

43 if atoms.has('initial_charges'): 

44 dct['initial_charges'] = atoms.get_initial_charges() 

45 if atoms.has('masses'): 

46 dct['masses'] = atoms.get_masses() 

47 if atoms.has('tags'): 

48 dct['tags'] = atoms.get_tags() 

49 if atoms.has('momenta'): 

50 dct['momenta'] = atoms.get_momenta() 

51 if atoms.constraints: 

52 dct['constraints'] = [c.todict() for c in atoms.constraints] 

53 if atoms.calc is not None: 

54 dct['calculator'] = atoms.calc.name.lower() 

55 dct['calculator_parameters'] = atoms.calc.todict() 

56 if len(atoms.calc.check_state(atoms)) == 0: 

57 for prop in all_properties: 

58 try: 

59 x = atoms.calc.get_property(prop, atoms, False) 

60 except PropertyNotImplementedError: 

61 pass 

62 else: 

63 if x is not None: 

64 dct[prop] = x 

65 return dct 

66 

67 

68class AtomsRow: 

69 mtime: float 

70 positions: np.ndarray 

71 id: int 

72 

73 def __init__(self, dct): 

74 if isinstance(dct, dict): 

75 dct = dct.copy() 

76 if 'calculator_parameters' in dct: 

77 # Earlier version of ASE would encode the calculator 

78 # parameter dict again and again and again ... 

79 while isinstance(dct['calculator_parameters'], str): 

80 dct['calculator_parameters'] = decode( 

81 dct['calculator_parameters']) 

82 else: 

83 dct = atoms2dict(dct) 

84 assert 'numbers' in dct 

85 self._constraints = dct.pop('constraints', []) 

86 self._constrained_forces = None 

87 self._data = dct.pop('data', {}) 

88 kvp = dct.pop('key_value_pairs', {}) 

89 self._keys = list(kvp.keys()) 

90 self.__dict__.update(kvp) 

91 self.__dict__.update(dct) 

92 if 'cell' not in dct: 

93 self.cell = np.zeros((3, 3)) 

94 if 'pbc' not in dct: 

95 self.pbc = np.zeros(3, bool) 

96 

97 def __contains__(self, key): 

98 return key in self.__dict__ 

99 

100 def __iter__(self): 

101 return (key for key in self.__dict__ if key[0] != '_') 

102 

103 def get(self, key, default=None): 

104 """Return value of key if present or default if not.""" 

105 return getattr(self, key, default) 

106 

107 @property 

108 def key_value_pairs(self): 

109 """Return dict of key-value pairs.""" 

110 return dict((key, self.get(key)) for key in self._keys) 

111 

112 def count_atoms(self): 

113 """Count atoms. 

114 

115 Return dict mapping chemical symbol strings to number of atoms. 

116 """ 

117 count = {} 

118 for symbol in self.symbols: 

119 count[symbol] = count.get(symbol, 0) + 1 

120 return count 

121 

122 def __getitem__(self, key): 

123 return getattr(self, key) 

124 

125 def __setitem__(self, key, value): 

126 setattr(self, key, value) 

127 

128 def __str__(self): 

129 return '<AtomsRow: formula={0}, keys={1}>'.format( 

130 self.formula, ','.join(self._keys)) 

131 

132 @property 

133 def constraints(self): 

134 """List of constraints.""" 

135 if not isinstance(self._constraints, list): 

136 # Lazy decoding: 

137 cs = decode(self._constraints) 

138 self._constraints = [] 

139 for c in cs: 

140 # Convert to new format: 

141 name = c.pop('__name__', None) 

142 if name: 

143 c = {'name': name, 'kwargs': c} 

144 if c['name'].startswith('ase'): 

145 c['name'] = c['name'].rsplit('.', 1)[1] 

146 self._constraints.append(c) 

147 return [dict2constraint(d) for d in self._constraints] 

148 

149 @property 

150 def data(self): 

151 """Data dict.""" 

152 if isinstance(self._data, str): 

153 self._data = decode(self._data) # lazy decoding 

154 elif isinstance(self._data, bytes): 

155 from ase.db.core import bytes_to_object 

156 self._data = bytes_to_object(self._data) # lazy decoding 

157 return FancyDict(self._data) 

158 

159 @property 

160 def natoms(self): 

161 """Number of atoms.""" 

162 return len(self.numbers) 

163 

164 @property 

165 def formula(self): 

166 """Chemical formula string.""" 

167 return Formula('', _tree=[(self.symbols, 1)]).format('metal') 

168 

169 @property 

170 def symbols(self): 

171 """List of chemical symbols.""" 

172 return [chemical_symbols[Z] for Z in self.numbers] 

173 

174 @property 

175 def fmax(self): 

176 """Maximum atomic force.""" 

177 forces = self.constrained_forces 

178 return (forces**2).sum(1).max()**0.5 

179 

180 @property 

181 def constrained_forces(self): 

182 """Forces after applying constraints.""" 

183 if self._constrained_forces is not None: 

184 return self._constrained_forces 

185 forces = self.forces 

186 constraints = self.constraints 

187 if constraints: 

188 forces = forces.copy() 

189 atoms = self.toatoms() 

190 for constraint in constraints: 

191 constraint.adjust_forces(atoms, forces) 

192 

193 self._constrained_forces = forces 

194 return forces 

195 

196 @property 

197 def smax(self): 

198 """Maximum stress tensor component.""" 

199 return (self.stress**2).max()**0.5 

200 

201 @property 

202 def mass(self): 

203 """Total mass.""" 

204 if 'masses' in self: 

205 return self.masses.sum() 

206 return atomic_masses[self.numbers].sum() 

207 

208 @property 

209 def volume(self): 

210 """Volume of unit cell.""" 

211 if self.cell is None: 

212 return None 

213 vol = abs(np.linalg.det(self.cell)) 

214 if vol == 0.0: 

215 raise AttributeError 

216 return vol 

217 

218 @property 

219 def charge(self): 

220 """Total charge.""" 

221 charges = self.get('initial_charges') 

222 if charges is None: 

223 return 0.0 

224 return charges.sum() 

225 

226 def toatoms(self, 

227 add_additional_information=False): 

228 """Create Atoms object.""" 

229 atoms = Atoms(self.numbers, 

230 self.positions, 

231 cell=self.cell, 

232 pbc=self.pbc, 

233 magmoms=self.get('initial_magmoms'), 

234 charges=self.get('initial_charges'), 

235 tags=self.get('tags'), 

236 masses=self.get('masses'), 

237 momenta=self.get('momenta'), 

238 constraint=self.constraints) 

239 

240 results = {} 

241 for prop in all_properties: 

242 if prop in self: 

243 results[prop] = self[prop] 

244 if results: 

245 atoms.calc = SinglePointCalculator(atoms, **results) 

246 atoms.calc.name = self.get('calculator', 'unknown') 

247 

248 if add_additional_information: 

249 atoms.info = {} 

250 atoms.info['unique_id'] = self.unique_id 

251 if self._keys: 

252 atoms.info['key_value_pairs'] = self.key_value_pairs 

253 data = self.get('data') 

254 if data: 

255 atoms.info['data'] = data 

256 

257 return atoms 

258 

259 

260def row2dct(row, key_descriptions) -> Dict[str, Any]: 

261 """Convert row to dict of things for printing or a web-page.""" 

262 

263 from ase.db.core import float_to_time_string, now 

264 

265 dct = {} 

266 

267 atoms = Atoms(cell=row.cell, pbc=row.pbc) 

268 dct['size'] = kptdensity2monkhorstpack(atoms, 

269 kptdensity=1.8, 

270 even=False) 

271 

272 dct['cell'] = [['{:.3f}'.format(a) for a in axis] for axis in row.cell] 

273 par = ['{:.3f}'.format(x) for x in cell_to_cellpar(row.cell)] 

274 dct['lengths'] = par[:3] 

275 dct['angles'] = par[3:] 

276 

277 stress = row.get('stress') 

278 if stress is not None: 

279 dct['stress'] = ', '.join('{0:.3f}'.format(s) for s in stress) 

280 

281 dct['formula'] = Formula(row.formula).format('abc') 

282 

283 dipole = row.get('dipole') 

284 if dipole is not None: 

285 dct['dipole'] = ', '.join('{0:.3f}'.format(d) for d in dipole) 

286 

287 data = row.get('data') 

288 if data: 

289 dct['data'] = ', '.join(data.keys()) 

290 

291 constraints = row.get('constraints') 

292 if constraints: 

293 dct['constraints'] = ', '.join(c.__class__.__name__ 

294 for c in constraints) 

295 

296 keys = ({'id', 'energy', 'fmax', 'smax', 'mass', 'age'} | 

297 set(key_descriptions) | 

298 set(row.key_value_pairs)) 

299 dct['table'] = [] 

300 

301 from ase.db.project import KeyDescription 

302 for key in keys: 

303 if key == 'age': 

304 age = float_to_time_string(now() - row.ctime, True) 

305 dct['table'].append(('ctime', 'Age', age)) 

306 continue 

307 value = row.get(key) 

308 if value is not None: 

309 if isinstance(value, float): 

310 value = '{:.3f}'.format(value) 

311 elif not isinstance(value, str): 

312 value = str(value) 

313 

314 nokeydesc = KeyDescription(key, '', '', '') 

315 keydesc = key_descriptions.get(key, nokeydesc) 

316 unit = keydesc.unit 

317 if unit: 

318 value += ' ' + unit 

319 dct['table'].append((key, keydesc.longdesc, value)) 

320 

321 return dct