Coverage for /builds/ase/ase/ase/collections/collection.py : 97.78%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import os.path as op
3from ase.db.row import AtomsRow
4from ase.io.jsonio import read_json
7class Collection:
8 """Collection of atomic configurations and associated data.
10 Example of use:
12 >>> from ase.collections import s22
13 >>> len(s22)
14 22
15 >>> s22.names[:3]
16 ['Ammonia_dimer', 'Water_dimer', 'Formic_acid_dimer']
17 >>> dimer = s22['Water_dimer']
18 >>> dimer.get_chemical_symbols()
19 ['O', 'H', 'H', 'O', 'H', 'H']
20 >>> s22.data['Ammonia_dimer']
21 {'cc_energy': -0.1375}
22 >>> sum(len(atoms) for atoms in s22)
23 414
24 """
25 def __init__(self, name):
26 """Create a collection lazily.
28 Will read data from json file when needed.
30 A collection can be iterated over to get the Atoms objects and indexed
31 with names to get individual members.
33 Attributes:
35 name: str
36 Name of collection.
37 data: dict
38 Data dictionary.
39 filename: str
40 Location of json file.
41 names: list
42 Names of configurations in the collection.
43 """
45 self.name = name
46 self._names = []
47 self._systems = {}
48 self._data = {}
49 self.filename = op.join(op.dirname(__file__), name + '.json')
51 def __getitem__(self, name):
52 self._read()
53 return self._systems[name].copy()
55 def has(self, name):
56 # Not __contains__() because __iter__ yields the systems.
57 self._read()
58 return name in self._systems
60 def __iter__(self):
61 for name in self.names:
62 yield self[name]
64 def __len__(self):
65 return len(self.names)
67 def __str__(self):
68 return '<{0}-collection, {1} systems: {2}, {3}, ...>'.format(
69 self.name, len(self), *self.names[:2])
71 def __repr__(self):
72 return 'Collection({0!r})'.format(self.name)
74 @property
75 def names(self):
76 self._read()
77 return list(self._names)
79 @property
80 def data(self):
81 self._read()
82 return self._data
84 def _read(self):
85 if self._names:
86 return
87 bigdct = read_json(self.filename)
88 for id in bigdct['ids']:
89 dct = bigdct[id]
90 kvp = dct['key_value_pairs']
91 name = str(kvp['name'])
92 self._names.append(name)
93 self._systems[name] = AtomsRow(dct).toatoms()
94 del kvp['name']
95 self._data[name] = dict((str(k), v) for k, v in kvp.items())