Coverage for /builds/ericyuan00000/ase/ase/outputs.py: 98.02%

1# fmt: off 

2 

3from abc import ABC, abstractmethod 

4from collections.abc import Mapping 

5from typing import Sequence, Union 

6 

7import numpy as np 

8 

9 

10class Properties(Mapping): 

11 def __init__(self, dct): 

12 self._dct = {} 

13 for name, value in dct.items(): 

14 self._setvalue(name, value) 

15 

16 def __len__(self): 

17 return len(self._dct) 

18 

19 def __iter__(self): 

20 return iter(self._dct) 

21 

22 def __getitem__(self, name): 

23 return self._dct[name] 

24 

25 def _setvalue(self, name, value): 

26 if name in self._dct: 

27 # Which error should we raise for already existing property? 

28 raise ValueError(f'{name} already set') 

29 

30 prop = all_outputs[name] 

31 value = prop.normalize_type(value) 

32 shape = np.shape(value) 

33 

34 if not self.shape_is_consistent(prop, value): 

35 raise ValueError(f'{name} has bad shape: {shape}') 

36 

37 for i, spec in enumerate(prop.shapespec): 

38 if not isinstance(spec, str) or spec in self._dct: 

39 continue 

40 self._setvalue(spec, shape[i]) 

41 

42 self._dct[name] = value 

43 

44 def shape_is_consistent(self, prop, value) -> bool: 

45 """Return whether shape of values is consistent with properties. 

46 

47 For example, forces of shape (7, 3) are consistent 

48 unless properties already have "natoms" with non-7 value. 

49 """ 

50 shapespec = prop.shapespec 

51 shape = np.shape(value) 

52 if len(shapespec) != len(shape): 

53 return False 

54 for dimspec, dim in zip(shapespec, shape): 

55 if isinstance(dimspec, str): 

56 dimspec = self._dct.get(dimspec, dim) 

57 if dimspec != dim: 

58 return False 

59 return True 

60 

61 def __repr__(self): 

62 clsname = type(self).__name__ 

63 return f'({clsname}({self._dct})' 

64 

65 

66all_outputs = {} 

67 

68 

69class Property(ABC): 

70 def __init__(self, name, dtype, shapespec): 

71 self.name = name 

72 assert dtype in [float, int] # Others? 

73 self.dtype = dtype 

74 self.shapespec = shapespec 

75 

76 @abstractmethod 

77 def normalize_type(self, value): 

78 ... 

79 

80 def __repr__(self) -> str: 

81 typename = self.dtype.__name__ # Extend to other than float/int? 

82 shape = ', '.join(str(dim) for dim in self.shapespec) 

83 return f'Property({self.name!r}, dtype={typename}, shape=[{shape}])' 

84 

85 

86class ScalarProperty(Property): 

87 def __init__(self, name, dtype): 

88 super().__init__(name, dtype, ()) 

89 

90 def normalize_type(self, value): 

91 if not np.isscalar(value): 

92 raise TypeError('Expected scalar') 

93 return self.dtype(value) 

94 

95 

96class ArrayProperty(Property): 

97 def normalize_type(self, value): 

98 if np.isscalar(value): 

99 raise TypeError('Expected array, got scalar') 

100 return np.asarray(value, dtype=self.dtype) 

101 

102 

103ShapeSpec = Union[str, int] 

104 

105 

106def _defineprop( 

107 name: str, 

108 dtype: type = float, 

109 shape: Union[ShapeSpec, Sequence[ShapeSpec]] = () 

110) -> Property: 

111 """Create, register, and return a property.""" 

112 

113 if isinstance(shape, (int, str)): 

114 shape = (shape,) 

115 

116 shape = tuple(shape) 

117 prop: Property 

118 if len(shape) == 0: 

119 prop = ScalarProperty(name, dtype) 

120 else: 

121 prop = ArrayProperty(name, dtype, shape) 

122 

123 assert name not in all_outputs, name 

124 all_outputs[name] = prop 

125 return prop 

126 

127 

128# Atoms, energy, forces, stress: 

129_defineprop('natoms', int) 

130_defineprop('energy', float) 

131_defineprop('energies', float, shape='natoms') 

132_defineprop('free_energy', float) 

133_defineprop('forces', float, shape=('natoms', 3)) 

134_defineprop('stress', float, shape=6) 

135_defineprop('stresses', float, shape=('natoms', 6)) 

136 

137# Electronic structure: 

138_defineprop('nbands', int) 

139_defineprop('nkpts', int) 

140_defineprop('nspins', int) 

141_defineprop('fermi_level', float) 

142_defineprop('kpoint_weights', float, shape='nkpts') 

143_defineprop('ibz_kpoints', float, shape=('nkpts', 3)) 

144_defineprop('eigenvalues', float, shape=('nspins', 'nkpts', 'nbands')) 

145_defineprop('occupations', float, shape=('nspins', 'nkpts', 'nbands')) 

146 

147# Miscellaneous: 

148_defineprop('dipole', float, shape=3) 

149_defineprop('charges', float, shape='natoms') 

150_defineprop('magmom', float) 

151_defineprop('magmoms', float, shape='natoms') # XXX spinors? 

152_defineprop('polarization', float, shape=3) 

153_defineprop('dielectric_tensor', float, shape=(3, 3)) 

154_defineprop('born_effective_charges', float, shape=('natoms', 3, 3)) 

155 

156# We might want to allow properties that are part of Atoms, such as 

157# positions, numbers, pbc, cell. It would be reasonable for those 

158# concepts to have a formalization outside the Atoms class. 

159 

160 

161# def to_singlepoint(self, atoms): 

162# from ase.calculators.singlepoint import SinglePointDFTCalculator 

163# return SinglePointDFTCalculator(atoms, 

164# efermi=self.fermi_level, 

165 

166# We can also retrieve (P)DOS and band structure. However: 

167# 

168# * Band structure may require bandpath, which is an input, and 

169# may not necessarily be easy or possible to reconstruct from 

170# the outputs. 

171# 

172# * Some calculators may produce the whole BandStructure object in 

173# one go (e.g. while parsing) 

174# 

175# * What about HOMO/LUMO? Can be obtained from 

176# eigenvalues/occupations, but some codes provide real data. We 

177# probably need to distinguish between HOMO/LUMO inferred by us 

178# versus values provided within the output. 

179# 

180# * HOMO is sometimes used as alternative reference energy for 

181# band structure. 

182# 

183# * What about spin-dependent (double) Fermi level? 

184# 

185# * What about 3D arrays? We will almost certainly want to be 

186# connected to an object that can load dynamically from a file.