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

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

47 

48 self.name = name 

49 self._names = [] 

50 self._systems = {} 

51 self._data = {} 

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

53 

54 def __getitem__(self, name): 

55 self._read() 

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

57 

58 def has(self, name): 

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

60 self._read() 

61 return name in self._systems 

62 

63 def __iter__(self): 

64 for name in self.names: 

65 yield self[name] 

66 

67 def __len__(self): 

68 return len(self.names) 

69 

70 def __str__(self): 

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

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

73 

74 def __repr__(self): 

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

76 

77 @property 

78 def names(self): 

79 self._read() 

80 return list(self._names) 

81 

82 @property 

83 def data(self): 

84 self._read() 

85 return self._data 

86 

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()}