Coverage for /builds/debichem-team/python-ase/ase/db/row.py: 92.76%

1from random import randint 

2from typing import Any, Dict 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.calculators.calculator import ( 

8 PropertyNotImplementedError, 

9 all_properties, 

10 kptdensity2monkhorstpack, 

11) 

12from ase.calculators.singlepoint import SinglePointCalculator 

13from ase.data import atomic_masses, chemical_symbols 

14from ase.formula import Formula 

15from ase.geometry import cell_to_cellpar 

16from ase.io.jsonio import decode 

17 

18 

19class FancyDict(dict): 

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

21 

22 def __getattr__(self, key): 

23 if key not in self: 

24 return dict.__getattribute__(self, key) 

25 value = self[key] 

26 if isinstance(value, dict): 

27 return FancyDict(value) 

28 return value 

29 

30 def __dir__(self): 

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

32 

33 

34def atoms2dict(atoms): 

35 dct = { 

36 'numbers': atoms.numbers, 

37 'positions': atoms.positions, 

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

39 if atoms.pbc.any(): 

40 dct['pbc'] = atoms.pbc 

41 if atoms.cell.any(): 

42 dct['cell'] = atoms.cell 

43 if atoms.has('initial_magmoms'): 

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

45 if atoms.has('initial_charges'): 

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

47 if atoms.has('masses'): 

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

49 if atoms.has('tags'): 

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

51 if atoms.has('momenta'): 

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

53 if atoms.constraints: 

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

55 if atoms.calc is not None: 

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

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

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

59 for prop in all_properties: 

60 try: 

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

62 except PropertyNotImplementedError: 

63 pass 

64 else: 

65 if x is not None: 

66 dct[prop] = x 

67 return dct 

68 

69 

70class AtomsRow: 

71 mtime: float 

72 positions: np.ndarray 

73 id: int 

74 

75 def __init__(self, dct): 

76 if isinstance(dct, dict): 

77 dct = dct.copy() 

78 if 'calculator_parameters' in dct: 

79 # Earlier version of ASE would encode the calculator 

80 # parameter dict again and again and again ... 

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

82 dct['calculator_parameters'] = decode( 

83 dct['calculator_parameters']) 

84 else: 

85 dct = atoms2dict(dct) 

86 assert 'numbers' in dct 

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

88 self._constrained_forces = None 

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

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

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

92 self.__dict__.update(kvp) 

93 self.__dict__.update(dct) 

94 if 'cell' not in dct: 

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

96 if 'pbc' not in dct: 

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

98 

99 def __contains__(self, key): 

100 return key in self.__dict__ 

101 

102 def __iter__(self): 

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

104 

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

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

107 return getattr(self, key, default) 

108 

109 @property 

110 def key_value_pairs(self): 

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

112 return {key: self.get(key) for key in self._keys} 

113 

114 def count_atoms(self): 

115 """Count atoms. 

116 

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

118 """ 

119 count = {} 

120 for symbol in self.symbols: 

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

122 return count 

123 

124 def __getitem__(self, key): 

125 return getattr(self, key) 

126 

127 def __setitem__(self, key, value): 

128 setattr(self, key, value) 

129 

130 def __str__(self): 

131 return '<AtomsRow: formula={}, keys={}>'.format( 

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

133 

134 @property 

135 def constraints(self): 

136 """List of constraints.""" 

137 from ase.constraints import dict2constraint 

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

139 # Lazy decoding: 

140 cs = decode(self._constraints) 

141 self._constraints = [] 

142 for c in cs: 

143 # Convert to new format: 

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

145 if name: 

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

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

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

149 self._constraints.append(c) 

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

151 

152 @property 

153 def data(self): 

154 """Data dict.""" 

155 if isinstance(self._data, str): 

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

157 elif isinstance(self._data, bytes): 

158 from ase.db.core import bytes_to_object 

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

160 return FancyDict(self._data) 

161 

162 @property 

163 def natoms(self): 

164 """Number of atoms.""" 

165 return len(self.numbers) 

166 

167 @property 

168 def formula(self): 

169 """Chemical formula string.""" 

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

171 

172 @property 

173 def symbols(self): 

174 """List of chemical symbols.""" 

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

176 

177 @property 

178 def fmax(self): 

179 """Maximum atomic force.""" 

180 forces = self.constrained_forces 

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

182 

183 @property 

184 def constrained_forces(self): 

185 """Forces after applying constraints.""" 

186 if self._constrained_forces is not None: 

187 return self._constrained_forces 

188 forces = self.forces 

189 constraints = self.constraints 

190 if constraints: 

191 forces = forces.copy() 

192 atoms = self.toatoms() 

193 for constraint in constraints: 

194 constraint.adjust_forces(atoms, forces) 

195 

196 self._constrained_forces = forces 

197 return forces 

198 

199 @property 

200 def smax(self): 

201 """Maximum stress tensor component.""" 

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

203 

204 @property 

205 def mass(self): 

206 """Total mass.""" 

207 if 'masses' in self: 

208 return self.masses.sum() 

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

210 

211 @property 

212 def volume(self): 

213 """Volume of unit cell.""" 

214 if self.cell is None: 

215 return None 

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

217 if vol == 0.0: 

218 raise AttributeError 

219 return vol 

220 

221 @property 

222 def charge(self): 

223 """Total charge.""" 

224 charges = self.get('initial_charges') 

225 if charges is None: 

226 return 0.0 

227 return charges.sum() 

228 

229 def toatoms(self, 

230 add_additional_information=False): 

231 """Create Atoms object.""" 

232 atoms = Atoms(self.numbers, 

233 self.positions, 

234 cell=self.cell, 

235 pbc=self.pbc, 

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

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

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

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

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

241 constraint=self.constraints) 

242 

243 results = {prop: self[prop] for prop in all_properties if prop in self} 

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'] = [[f'{a:.3f}' for a in axis] for axis in row.cell] 

273 par = [f'{x:.3f}' 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(f'{s:.3f}' 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(f'{d:.3f}' 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 = f'{value:.3f}' 

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