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

1# fmt: off 

2 

3import os.path as op 

4 

5from ase.db.row import AtomsRow 

6from ase.io.jsonio import read_json 

7 

8 

9class Collection: 

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

11 

12 Example of use: 

13 

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 """ 

27 

28 def __init__(self, name): 

29 """Create a collection lazily. 

30 

31 Will read data from json file when needed. 

32 

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

34 with names to get individual members. 

35 

36 Attributes 

37 ---------- 

38 

39 name: str 

40 Name of collection. 

41 data: dict 

42 Data dictionary. 

43 filename: str 

44 Location of json file. 

45 names: list 

46 Names of configurations in the collection. 

47 """ 

48 

49 self.name = name 

50 self._names = [] 

51 self._systems = {} 

52 self._data = {} 

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

54 

55 def __getitem__(self, name): 

56 self._read() 

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

58 

59 def has(self, name): 

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

61 self._read() 

62 return name in self._systems 

63 

64 def __iter__(self): 

65 for name in self.names: 

66 yield self[name] 

67 

68 def __len__(self): 

69 return len(self.names) 

70 

71 def __str__(self): 

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

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

74 

75 def __repr__(self): 

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

77 

78 @property 

79 def names(self): 

80 self._read() 

81 return list(self._names) 

82 

83 @property 

84 def data(self): 

85 self._read() 

86 return self._data 

87 

88 def _read(self): 

89 if self._names: 

90 return 

91 bigdct = read_json(self.filename) 

92 for id in bigdct['ids']: 

93 dct = bigdct[id] 

94 kvp = dct['key_value_pairs'] 

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

96 self._names.append(name) 

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

98 del kvp['name'] 

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