Hide keyboard shortcuts

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 

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 def __init__(self, name): 

26 """Create a collection lazily. 

27 

28 Will read data from json file when needed. 

29 

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

31 with names to get individual members. 

32 

33 Attributes: 

34 

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

44 

45 self.name = name 

46 self._names = [] 

47 self._systems = {} 

48 self._data = {} 

49 self.filename = op.join(op.dirname(__file__), name + '.json') 

50 

51 def __getitem__(self, name): 

52 self._read() 

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

54 

55 def has(self, name): 

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

57 self._read() 

58 return name in self._systems 

59 

60 def __iter__(self): 

61 for name in self.names: 

62 yield self[name] 

63 

64 def __len__(self): 

65 return len(self.names) 

66 

67 def __str__(self): 

68 return '<{0}-collection, {1} systems: {2}, {3}, ...>'.format( 

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

70 

71 def __repr__(self): 

72 return 'Collection({0!r})'.format(self.name) 

73 

74 @property 

75 def names(self): 

76 self._read() 

77 return list(self._names) 

78 

79 @property 

80 def data(self): 

81 self._read() 

82 return self._data 

83 

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