Coverage for /builds/ase/ase/ase/db/row.py: 92.76%
221 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3import uuid
4from typing import Any, Dict
6import numpy as np
8from ase import Atoms
9from ase.calculators.calculator import (
10 PropertyNotImplementedError,
11 all_properties,
12 kptdensity2monkhorstpack,
13)
14from ase.calculators.singlepoint import SinglePointCalculator
15from ase.data import atomic_masses, chemical_symbols
16from ase.formula import Formula
17from ase.geometry import cell_to_cellpar
18from ase.io.jsonio import decode
21class FancyDict(dict):
22 """Dictionary with keys available as attributes also."""
24 def __getattr__(self, key):
25 if key not in self:
26 return dict.__getattribute__(self, key)
27 value = self[key]
28 if isinstance(value, dict):
29 return FancyDict(value)
30 return value
32 def __dir__(self):
33 return self.keys() # for tab-completion
36def atoms2dict(atoms):
37 dct = {
38 'numbers': atoms.numbers,
39 'positions': atoms.positions,
40 'unique_id': uuid.uuid4().hex}
41 if atoms.pbc.any():
42 dct['pbc'] = atoms.pbc
43 if atoms.cell.any():
44 dct['cell'] = atoms.cell
45 if atoms.has('initial_magmoms'):
46 dct['initial_magmoms'] = atoms.get_initial_magnetic_moments()
47 if atoms.has('initial_charges'):
48 dct['initial_charges'] = atoms.get_initial_charges()
49 if atoms.has('masses'):
50 dct['masses'] = atoms.get_masses()
51 if atoms.has('tags'):
52 dct['tags'] = atoms.get_tags()
53 if atoms.has('momenta'):
54 dct['momenta'] = atoms.get_momenta()
55 if atoms.constraints:
56 dct['constraints'] = [c.todict() for c in atoms.constraints]
57 if atoms.calc is not None:
58 dct['calculator'] = atoms.calc.name.lower()
59 dct['calculator_parameters'] = atoms.calc.todict()
60 if len(atoms.calc.check_state(atoms)) == 0:
61 for prop in all_properties:
62 try:
63 x = atoms.calc.get_property(prop, atoms, False)
64 except PropertyNotImplementedError:
65 pass
66 else:
67 if x is not None:
68 dct[prop] = x
69 return dct
72class AtomsRow:
73 mtime: float
74 positions: np.ndarray
75 id: int
77 def __init__(self, dct):
78 if isinstance(dct, dict):
79 dct = dct.copy()
80 if 'calculator_parameters' in dct:
81 # Earlier version of ASE would encode the calculator
82 # parameter dict again and again and again ...
83 while isinstance(dct['calculator_parameters'], str):
84 dct['calculator_parameters'] = decode(
85 dct['calculator_parameters'])
86 else:
87 dct = atoms2dict(dct)
88 assert 'numbers' in dct
89 self._constraints = dct.pop('constraints', [])
90 self._constrained_forces = None
91 self._data = dct.pop('data', {})
92 kvp = dct.pop('key_value_pairs', {})
93 self._keys = list(kvp.keys())
94 self.__dict__.update(kvp)
95 self.__dict__.update(dct)
96 if 'cell' not in dct:
97 self.cell = np.zeros((3, 3))
98 if 'pbc' not in dct:
99 self.pbc = np.zeros(3, bool)
101 def __contains__(self, key):
102 return key in self.__dict__
104 def __iter__(self):
105 return (key for key in self.__dict__ if key[0] != '_')
107 def get(self, key, default=None):
108 """Return value of key if present or default if not."""
109 return getattr(self, key, default)
111 @property
112 def key_value_pairs(self):
113 """Return dict of key-value pairs."""
114 return {key: self.get(key) for key in self._keys}
116 def count_atoms(self):
117 """Count atoms.
119 Return dict mapping chemical symbol strings to number of atoms.
120 """
121 count = {}
122 for symbol in self.symbols:
123 count[symbol] = count.get(symbol, 0) + 1
124 return count
126 def __getitem__(self, key):
127 return getattr(self, key)
129 def __setitem__(self, key, value):
130 setattr(self, key, value)
132 def __str__(self):
133 return '<AtomsRow: formula={}, keys={}>'.format(
134 self.formula, ','.join(self._keys))
136 @property
137 def constraints(self):
138 """List of constraints."""
139 from ase.constraints import dict2constraint
140 if not isinstance(self._constraints, list):
141 # Lazy decoding:
142 cs = decode(self._constraints)
143 self._constraints = []
144 for c in cs:
145 # Convert to new format:
146 name = c.pop('__name__', None)
147 if name:
148 c = {'name': name, 'kwargs': c}
149 if c['name'].startswith('ase'):
150 c['name'] = c['name'].rsplit('.', 1)[1]
151 self._constraints.append(c)
152 return [dict2constraint(d) for d in self._constraints]
154 @property
155 def data(self):
156 """Data dict."""
157 if isinstance(self._data, str):
158 self._data = decode(self._data) # lazy decoding
159 elif isinstance(self._data, bytes):
160 from ase.db.core import bytes_to_object
161 self._data = bytes_to_object(self._data) # lazy decoding
162 return FancyDict(self._data)
164 @property
165 def natoms(self):
166 """Number of atoms."""
167 return len(self.numbers)
169 @property
170 def formula(self):
171 """Chemical formula string."""
172 return Formula('', _tree=[(self.symbols, 1)]).format('metal')
174 @property
175 def symbols(self):
176 """List of chemical symbols."""
177 return [chemical_symbols[Z] for Z in self.numbers]
179 @property
180 def fmax(self):
181 """Maximum atomic force."""
182 forces = self.constrained_forces
183 return (forces**2).sum(1).max()**0.5
185 @property
186 def constrained_forces(self):
187 """Forces after applying constraints."""
188 if self._constrained_forces is not None:
189 return self._constrained_forces
190 forces = self.forces
191 constraints = self.constraints
192 if constraints:
193 forces = forces.copy()
194 atoms = self.toatoms()
195 for constraint in constraints:
196 constraint.adjust_forces(atoms, forces)
198 self._constrained_forces = forces
199 return forces
201 @property
202 def smax(self):
203 """Maximum stress tensor component."""
204 return (self.stress**2).max()**0.5
206 @property
207 def mass(self):
208 """Total mass."""
209 if 'masses' in self:
210 return self.masses.sum()
211 return atomic_masses[self.numbers].sum()
213 @property
214 def volume(self):
215 """Volume of unit cell."""
216 if self.cell is None:
217 return None
218 vol = abs(np.linalg.det(self.cell))
219 if vol == 0.0:
220 raise AttributeError
221 return vol
223 @property
224 def charge(self):
225 """Total charge."""
226 charges = self.get('initial_charges')
227 if charges is None:
228 return 0.0
229 return charges.sum()
231 def toatoms(self,
232 add_additional_information=False):
233 """Create Atoms object."""
234 atoms = Atoms(self.numbers,
235 self.positions,
236 cell=self.cell,
237 pbc=self.pbc,
238 magmoms=self.get('initial_magmoms'),
239 charges=self.get('initial_charges'),
240 tags=self.get('tags'),
241 masses=self.get('masses'),
242 momenta=self.get('momenta'),
243 constraint=self.constraints)
245 results = {prop: self[prop] for prop in all_properties if prop in self}
246 if results:
247 atoms.calc = SinglePointCalculator(atoms, **results)
248 atoms.calc.name = self.get('calculator', 'unknown')
250 if add_additional_information:
251 atoms.info = {}
252 atoms.info['unique_id'] = self.unique_id
253 if self._keys:
254 atoms.info['key_value_pairs'] = self.key_value_pairs
255 data = self.get('data')
256 if data:
257 atoms.info['data'] = data
259 return atoms
262def row2dct(row, key_descriptions) -> Dict[str, Any]:
263 """Convert row to dict of things for printing or a web-page."""
265 from ase.db.core import float_to_time_string, now
267 dct = {}
269 atoms = Atoms(cell=row.cell, pbc=row.pbc)
270 dct['size'] = kptdensity2monkhorstpack(atoms,
271 kptdensity=1.8,
272 even=False)
274 dct['cell'] = [[f'{a:.3f}' for a in axis] for axis in row.cell]
275 par = [f'{x:.3f}' for x in cell_to_cellpar(row.cell)]
276 dct['lengths'] = par[:3]
277 dct['angles'] = par[3:]
279 stress = row.get('stress')
280 if stress is not None:
281 dct['stress'] = ', '.join(f'{s:.3f}' for s in stress)
283 dct['formula'] = Formula(row.formula).format('abc')
285 dipole = row.get('dipole')
286 if dipole is not None:
287 dct['dipole'] = ', '.join(f'{d:.3f}' for d in dipole)
289 data = row.get('data')
290 if data:
291 dct['data'] = ', '.join(data.keys())
293 constraints = row.get('constraints')
294 if constraints:
295 dct['constraints'] = ', '.join(c.__class__.__name__
296 for c in constraints)
298 keys = ({'id', 'energy', 'fmax', 'smax', 'mass', 'age'} |
299 set(key_descriptions) |
300 set(row.key_value_pairs))
301 dct['table'] = []
303 from ase.db.project import KeyDescription
304 for key in keys:
305 if key == 'age':
306 age = float_to_time_string(now() - row.ctime, True)
307 dct['table'].append(('ctime', 'Age', age))
308 continue
309 value = row.get(key)
310 if value is not None:
311 if isinstance(value, float):
312 value = f'{value:.3f}'
313 elif not isinstance(value, str):
314 value = str(value)
316 nokeydesc = KeyDescription(key, '', '', '')
317 keydesc = key_descriptions.get(key, nokeydesc)
318 unit = keydesc.unit
319 if unit:
320 value += ' ' + unit
321 dct['table'].append((key, keydesc.longdesc, value))
323 return dct