Coverage for /builds/ase/ase/ase/collections/collection.py: 97.78%
45 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3import os.path as op
5from ase.db.row import AtomsRow
6from ase.io.jsonio import read_json
9class Collection:
10 """Collection of atomic configurations and associated data.
12 Example of use:
14 >>> from ase.collections import s22
15 >>> len(s22)
16 22
17 >>> s22.names[:3]
18 ['Ammonia_dimer', 'Water_dimer', 'Formic_acid_dimer']
19 >>> dimer = s22['Water_dimer']
20 >>> dimer.get_chemical_symbols()
21 ['O', 'H', 'H', 'O', 'H', 'H']
22 >>> s22.data['Ammonia_dimer']
23 {'cc_energy': -0.1375}
24 >>> sum(len(atoms) for atoms in s22)
25 414
26 """
28 def __init__(self, name):
29 """Create a collection lazily.
31 Will read data from json file when needed.
33 A collection can be iterated over to get the Atoms objects and indexed
34 with names to get individual members.
36 Attributes:
38 name: str
39 Name of collection.
40 data: dict
41 Data dictionary.
42 filename: str
43 Location of json file.
44 names: list
45 Names of configurations in the collection.
46 """
48 self.name = name
49 self._names = []
50 self._systems = {}
51 self._data = {}
52 self.filename = op.join(op.dirname(__file__), f'{name}.json')
54 def __getitem__(self, name):
55 self._read()
56 return self._systems[name].copy()
58 def has(self, name):
59 # Not __contains__() because __iter__ yields the systems.
60 self._read()
61 return name in self._systems
63 def __iter__(self):
64 for name in self.names:
65 yield self[name]
67 def __len__(self):
68 return len(self.names)
70 def __str__(self):
71 return '<{}-collection, {} systems: {}, {}, ...>'.format(
72 self.name, len(self), *self.names[:2])
74 def __repr__(self):
75 return f'Collection({self.name!r})'
77 @property
78 def names(self):
79 self._read()
80 return list(self._names)
82 @property
83 def data(self):
84 self._read()
85 return self._data
87 def _read(self):
88 if self._names:
89 return
90 bigdct = read_json(self.filename)
91 for id in bigdct['ids']:
92 dct = bigdct[id]
93 kvp = dct['key_value_pairs']
94 name = str(kvp['name'])
95 self._names.append(name)
96 self._systems[name] = AtomsRow(dct).toatoms()
97 del kvp['name']
98 self._data[name] = {str(k): v for k, v in kvp.items()}