1from abc import ABC, abstractmethod 

2from collections.abc import Mapping 

3from typing import Sequence, Union 

4 

5import numpy as np 

6 

7 

8class Properties(Mapping): 

9 def __init__(self, dct): 

10 self._dct = {} 

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

12 self._setvalue(name, value) 

13 

14 def __len__(self): 

15 return len(self._dct) 

16 

17 def __iter__(self): 

18 return iter(self._dct) 

19 

20 def __getitem__(self, name): 

21 return self._dct[name] 

22 

23 def _setvalue(self, name, value): 

24 if name in self._dct: 

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

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

27 

28 prop = all_outputs[name] 

29 value = prop.normalize_type(value) 

30 shape = np.shape(value) 

31 

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

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

34 

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

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

37 continue 

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

39 

40 self._dct[name] = value 

41 

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

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

44 

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

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

47 """ 

48 shapespec = prop.shapespec 

49 shape = np.shape(value) 

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

51 return False 

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

53 if isinstance(dimspec, str): 

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

55 if dimspec != dim: 

56 return False 

57 return True 

58 

59 def __repr__(self): 

60 clsname = type(self).__name__ 

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

62 

63 

64all_outputs = {} 

65 

66 

67class Property(ABC): 

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

69 self.name = name 

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

71 self.dtype = dtype 

72 self.shapespec = shapespec 

73 

74 @abstractmethod 

75 def normalize_type(self, value): 

76 ... 

77 

78 def __repr__(self) -> str: 

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

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

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

82 

83 

84class ScalarProperty(Property): 

85 def __init__(self, name, dtype): 

86 super().__init__(name, dtype, tuple()) 

87 

88 def normalize_type(self, value): 

89 if not np.isscalar(value): 

90 raise TypeError('Expected scalar') 

91 return self.dtype(value) 

92 

93 

94class ArrayProperty(Property): 

95 def normalize_type(self, value): 

96 if np.isscalar(value): 

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

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

99 

100 

101ShapeSpec = Union[str, int] 

102 

103 

104def _defineprop( 

105 name: str, 

106 dtype: type = float, 

107 shape: Union[ShapeSpec, Sequence[ShapeSpec]] = tuple() 

108) -> Property: 

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

110 

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

112 shape = (shape,) 

113 

114 shape = tuple(shape) 

115 prop: Property 

116 if len(shape) == 0: 

117 prop = ScalarProperty(name, dtype) 

118 else: 

119 prop = ArrayProperty(name, dtype, shape) 

120 all_outputs[name] = prop 

121 return prop 

122 

123 

124# Atoms, energy, forces, stress: 

125_defineprop('natoms', int) 

126_defineprop('energy', float) 

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

128_defineprop('free_energy', float) 

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

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

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

132 

133# Electronic structure: 

134_defineprop('nbands', int) 

135_defineprop('nkpts', int) 

136_defineprop('nspins', int) 

137_defineprop('fermi_level', float) 

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

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

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

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

142 

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

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

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

146 

147 

148#def to_singlepoint(self, atoms): 

149# from ase.calculators.singlepoint import SinglePointDFTCalculator 

150# return SinglePointDFTCalculator(atoms, 

151# efermi=self.fermi_level, 

152 

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

154# 

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

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

157# the outputs. 

158# 

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

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

161# 

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

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

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

165# versus values provided within the output. 

166# 

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

168# band structure. 

169# 

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

171# 

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

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