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

1import os.path as op 

2 

3from ase.db.row import AtomsRow 

4from ase.io.jsonio import read_json 

5 

6 

7class Collection: 

8 """Collection of atomic configurations and associated data. 

9 

10 Example of use: 

11 

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 

26 def __init__(self, name): 

27 """Create a collection lazily. 

28 

29 Will read data from json file when needed. 

30 

31 A collection can be iterated over to get the Atoms objects and indexed 

32 with names to get individual members. 

33 

34 Attributes: 

35 

36 name: str 

37 Name of collection. 

38 data: dict 

39 Data dictionary. 

40 filename: str 

41 Location of json file. 

42 names: list 

43 Names of configurations in the collection. 

44 """ 

45 

46 self.name = name 

47 self._names = [] 

48 self._systems = {} 

49 self._data = {} 

50 self.filename = op.join(op.dirname(__file__), f'{name}.json') 

51 

52 def __getitem__(self, name): 

53 self._read() 

54 return self._systems[name].copy() 

55 

56 def has(self, name): 

57 # Not __contains__() because __iter__ yields the systems. 

58 self._read() 

59 return name in self._systems 

60 

61 def __iter__(self): 

62 for name in self.names: 

63 yield self[name] 

64 

65 def __len__(self): 

66 return len(self.names) 

67 

68 def __str__(self): 

69 return '<{}-collection, {} systems: {}, {}, ...>'.format( 

70 self.name, len(self), *self.names[:2]) 

71 

72 def __repr__(self): 

73 return f'Collection({self.name!r})' 

74 

75 @property 

76 def names(self): 

77 self._read() 

78 return list(self._names) 

79 

80 @property 

81 def data(self): 

82 self._read() 

83 return self._data 

84 

85 def _read(self): 

86 if self._names: 

87 return 

88 bigdct = read_json(self.filename) 

89 for id in bigdct['ids']: 

90 dct = bigdct[id] 

91 kvp = dct['key_value_pairs'] 

92 name = str(kvp['name']) 

93 self._names.append(name) 

94 self._systems[name] = AtomsRow(dct).toatoms() 

95 del kvp['name'] 

96 self._data[name] = {str(k): v for k, v in kvp.items()}