Coverage for /builds/kinetik161/ase/ase/db/row.py: 92.86%
224 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1from random import randint
2from typing import Any, Dict
4import numpy as np
6from ase import Atoms
7from ase.calculators.calculator import (PropertyNotImplementedError,
8 all_properties,
9 kptdensity2monkhorstpack)
10from ase.calculators.singlepoint import SinglePointCalculator
11from ase.data import atomic_masses, chemical_symbols
12from ase.formula import Formula
13from ase.geometry import cell_to_cellpar
14from ase.io.jsonio import decode
17class FancyDict(dict):
18 """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
28 def __dir__(self):
29 return self.keys() # for tab-completion
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
68class AtomsRow:
69 mtime: float
70 positions: np.ndarray
71 id: int
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)
97 def __contains__(self, key):
98 return key in self.__dict__
100 def __iter__(self):
101 return (key for key in self.__dict__ if key[0] != '_')
103 def get(self, key, default=None):
104 """Return value of key if present or default if not."""
105 return getattr(self, key, default)
107 @property
108 def key_value_pairs(self):
109 """Return dict of key-value pairs."""
110 return {key: self.get(key) for key in self._keys}
112 def count_atoms(self):
113 """Count atoms.
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
122 def __getitem__(self, key):
123 return getattr(self, key)
125 def __setitem__(self, key, value):
126 setattr(self, key, value)
128 def __str__(self):
129 return '<AtomsRow: formula={}, keys={}>'.format(
130 self.formula, ','.join(self._keys))
132 @property
133 def constraints(self):
134 """List of constraints."""
135 from ase.constraints import dict2constraint
136 if not isinstance(self._constraints, list):
137 # Lazy decoding:
138 cs = decode(self._constraints)
139 self._constraints = []
140 for c in cs:
141 # Convert to new format:
142 name = c.pop('__name__', None)
143 if name:
144 c = {'name': name, 'kwargs': c}
145 if c['name'].startswith('ase'):
146 c['name'] = c['name'].rsplit('.', 1)[1]
147 self._constraints.append(c)
148 return [dict2constraint(d) for d in self._constraints]
150 @property
151 def data(self):
152 """Data dict."""
153 if isinstance(self._data, str):
154 self._data = decode(self._data) # lazy decoding
155 elif isinstance(self._data, bytes):
156 from ase.db.core import bytes_to_object
157 self._data = bytes_to_object(self._data) # lazy decoding
158 return FancyDict(self._data)
160 @property
161 def natoms(self):
162 """Number of atoms."""
163 return len(self.numbers)
165 @property
166 def formula(self):
167 """Chemical formula string."""
168 return Formula('', _tree=[(self.symbols, 1)]).format('metal')
170 @property
171 def symbols(self):
172 """List of chemical symbols."""
173 return [chemical_symbols[Z] for Z in self.numbers]
175 @property
176 def fmax(self):
177 """Maximum atomic force."""
178 forces = self.constrained_forces
179 return (forces**2).sum(1).max()**0.5
181 @property
182 def constrained_forces(self):
183 """Forces after applying constraints."""
184 if self._constrained_forces is not None:
185 return self._constrained_forces
186 forces = self.forces
187 constraints = self.constraints
188 if constraints:
189 forces = forces.copy()
190 atoms = self.toatoms()
191 for constraint in constraints:
192 constraint.adjust_forces(atoms, forces)
194 self._constrained_forces = forces
195 return forces
197 @property
198 def smax(self):
199 """Maximum stress tensor component."""
200 return (self.stress**2).max()**0.5
202 @property
203 def mass(self):
204 """Total mass."""
205 if 'masses' in self:
206 return self.masses.sum()
207 return atomic_masses[self.numbers].sum()
209 @property
210 def volume(self):
211 """Volume of unit cell."""
212 if self.cell is None:
213 return None
214 vol = abs(np.linalg.det(self.cell))
215 if vol == 0.0:
216 raise AttributeError
217 return vol
219 @property
220 def charge(self):
221 """Total charge."""
222 charges = self.get('initial_charges')
223 if charges is None:
224 return 0.0
225 return charges.sum()
227 def toatoms(self,
228 add_additional_information=False):
229 """Create Atoms object."""
230 atoms = Atoms(self.numbers,
231 self.positions,
232 cell=self.cell,
233 pbc=self.pbc,
234 magmoms=self.get('initial_magmoms'),
235 charges=self.get('initial_charges'),
236 tags=self.get('tags'),
237 masses=self.get('masses'),
238 momenta=self.get('momenta'),
239 constraint=self.constraints)
241 results = {}
242 for prop in all_properties:
243 if prop in self:
244 results[prop] = self[prop]
245 if results:
246 atoms.calc = SinglePointCalculator(atoms, **results)
247 atoms.calc.name = self.get('calculator', 'unknown')
249 if add_additional_information:
250 atoms.info = {}
251 atoms.info['unique_id'] = self.unique_id
252 if self._keys:
253 atoms.info['key_value_pairs'] = self.key_value_pairs
254 data = self.get('data')
255 if data:
256 atoms.info['data'] = data
258 return atoms
261def row2dct(row, key_descriptions) -> Dict[str, Any]:
262 """Convert row to dict of things for printing or a web-page."""
264 from ase.db.core import float_to_time_string, now
266 dct = {}
268 atoms = Atoms(cell=row.cell, pbc=row.pbc)
269 dct['size'] = kptdensity2monkhorstpack(atoms,
270 kptdensity=1.8,
271 even=False)
273 dct['cell'] = [[f'{a:.3f}' for a in axis] for axis in row.cell]
274 par = [f'{x:.3f}' for x in cell_to_cellpar(row.cell)]
275 dct['lengths'] = par[:3]
276 dct['angles'] = par[3:]
278 stress = row.get('stress')
279 if stress is not None:
280 dct['stress'] = ', '.join(f'{s:.3f}' for s in stress)
282 dct['formula'] = Formula(row.formula).format('abc')
284 dipole = row.get('dipole')
285 if dipole is not None:
286 dct['dipole'] = ', '.join(f'{d:.3f}' for d in dipole)
288 data = row.get('data')
289 if data:
290 dct['data'] = ', '.join(data.keys())
292 constraints = row.get('constraints')
293 if constraints:
294 dct['constraints'] = ', '.join(c.__class__.__name__
295 for c in constraints)
297 keys = ({'id', 'energy', 'fmax', 'smax', 'mass', 'age'} |
298 set(key_descriptions) |
299 set(row.key_value_pairs))
300 dct['table'] = []
302 from ase.db.project import KeyDescription
303 for key in keys:
304 if key == 'age':
305 age = float_to_time_string(now() - row.ctime, True)
306 dct['table'].append(('ctime', 'Age', age))
307 continue
308 value = row.get(key)
309 if value is not None:
310 if isinstance(value, float):
311 value = f'{value:.3f}'
312 elif not isinstance(value, str):
313 value = str(value)
315 nokeydesc = KeyDescription(key, '', '', '')
316 keydesc = key_descriptions.get(key, nokeydesc)
317 unit = keydesc.unit
318 if unit:
319 value += ' ' + unit
320 dct['table'].append((key, keydesc.longdesc, value))
322 return dct