Source code for ase.collections.collection

import os.path as op

from ase.db.row import AtomsRow
from ase.io.jsonio import read_json


[docs]class Collection: """Collection of atomic configurations and associated data. Example of use: >>> from ase.collections import s22 >>> len(s22) 22 >>> s22.names[:3] ['Ammonia_dimer', 'Water_dimer', 'Formic_acid_dimer'] >>> dimer = s22['Water_dimer'] >>> dimer.get_chemical_symbols() ['O', 'H', 'H', 'O', 'H', 'H'] >>> s22.data['Ammonia_dimer'] {'cc_energy': -0.1375} >>> sum(len(atoms) for atoms in s22) 414 """ def __init__(self, name): """Create a collection lazily. Will read data from json file when needed. A collection can be iterated over to get the Atoms objects and indexed with names to get individual members. Attributes: name: str Name of collection. data: dict Data dictionary. filename: str Location of json file. names: list Names of configurations in the collection. """ self.name = name self._names = [] self._systems = {} self._data = {} self.filename = op.join(op.dirname(__file__), f'{name}.json') def __getitem__(self, name): self._read() return self._systems[name].copy() def has(self, name): # Not __contains__() because __iter__ yields the systems. self._read() return name in self._systems def __iter__(self): for name in self.names: yield self[name] def __len__(self): return len(self.names) def __str__(self): return '<{}-collection, {} systems: {}, {}, ...>'.format( self.name, len(self), *self.names[:2]) def __repr__(self): return f'Collection({self.name!r})' @property def names(self): self._read() return list(self._names) @property def data(self): self._read() return self._data def _read(self): if self._names: return bigdct = read_json(self.filename) for id in bigdct['ids']: dct = bigdct[id] kvp = dct['key_value_pairs'] name = str(kvp['name']) self._names.append(name) self._systems[name] = AtomsRow(dct).toatoms() del kvp['name'] self._data[name] = {str(k): v for k, v in kvp.items()}